{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "962be290",
   "metadata": {},
   "source": [
    "# Transformations\n",
    "\n",
    "In general, JAX transformations (transforms) operate on [pytrees](https://jax.readthedocs.io/en/latest/pytrees.html) of `jax.Array`s\n",
    "and abide by value semantics. This presents a challenge for Flax NNX, which represents `nnx.Module`s as regular Python objects\n",
    "that follow reference semantics. To address this, Flax NNX introduced its own set of transforms that extend JAX\n",
    "transforms to allow `nnx.Module`s and other Flax NNX objects to be passed in and out of transforms while preserving\n",
    "reference semantics.\n",
    "\n",
    "Flax NNX transforms should feel quite familiar if you have used JAX transforms before. They use the\n",
    "same APIs and behave like the JAX transforms when only working with pytrees of `jax.Array`s. However, when working with\n",
    "Flax NNX objects, they allow Python's reference semantics to be preserved for these objects, this includes:\n",
    "\n",
    "* Preserving shared references across multiple objects in the inputs and outputs of the transformation.\n",
    "* Propagating any state changes made to the objects inside the transformation to the objects outside the transformation.\n",
    "* Enforcing consistency of how objects are transformed when aliases are present across multiple inputs and outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d645146",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "from jax import numpy as jnp, random\n",
    "from flax import nnx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b44fb248",
   "metadata": {},
   "source": [
    "Throughout this guide, `nnx.vmap` is used as a case study to demonstrate how Flax NNX transforms work. However, the principles\n",
    "outlined in this document extends to all transforms.\n",
    "\n",
    "## Basic example\n",
    "\n",
    "To begin, let's look at a simple example of using `nnx.vmap` to extend an element wise `vector_dot` function to work on\n",
    "batched inputs. We will define a `Weights` Module with no methods to hold some parameters, these weights will be passed\n",
    "as an input to the `vector_dot` function along with some data. Both the weights and data will be batched on axis `0` and we will use\n",
    "`nnx.vmap` to apply `vector_dot` to each batch element, and the result will be a batched on axis `1`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4eab27a4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y.shape = (3, 10)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/ivyzheng/envs/f1/lib/python3.12/site-packages/treescope/renderers.py:251: UserWarning: Ignoring error inside wrapper hook <function use_autovisualizer_if_present at 0x11c70dda0>:\n",
      "Traceback (most recent call last):\n",
      "  File \"/Users/ivyzheng/envs/f1/lib/python3.12/site-packages/treescope/renderers.py\", line 225, in _render_subtree\n",
      "    postprocessed_result = hook(\n",
      "                           ^^^^^\n",
      "  File \"/Users/ivyzheng/envs/f1/lib/python3.12/site-packages/treescope/_internal/handlers/autovisualizer_hook.py\", line 47, in use_autovisualizer_if_present\n",
      "    result = autoviz(node, path)\n",
      "             ^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/Users/ivyzheng/envs/f1/lib/python3.12/site-packages/treescope/_internal/api/array_autovisualizer.py\", line 306, in __call__\n",
      "    jax.sharding.PositionalSharding\n",
      "  File \"/Users/ivyzheng/envs/f1/lib/python3.12/site-packages/jax/_src/deprecations.py\", line 54, in getattr\n",
      "    raise AttributeError(message)\n",
      "AttributeError: jax.sharding.PositionalSharding was deprecated in JAX v0.6.0 and removed in JAX v0.7.0\n",
      "\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<script> (()=>{ if (customElements.get('treescope-container') === undefined) { class TreescopeContainer extends HTMLElement { constructor() { super(); this.attachShadow({mode: \"open\"}); this.defns = {}; this.state = {}; } } customElements.define(\"treescope-container\", TreescopeContainer); } if (customElements.get('treescope-run-here') === undefined) { class RunHere extends HTMLElement { constructor() { super() } connectedCallback() { const run = child => { const fn = new Function(child.textContent); child.textContent = \"\"; fn.call(this); this.remove(); }; const child = this.querySelector(\"script\"); if (child) { run(child); } else { new MutationObserver(()=>{ run(this.querySelector(\"script\")); }).observe(this, {childList: true}); } } } customElements.define(\"treescope-run-here\", RunHere); } })(); </script> <treescope-container class=\"treescope_out_5b83677f05674a588929eaebe4d0dfc0\" style=\"display:block\"></treescope-container> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_5b83677f05674a588929eaebe4d0dfc0\")) .filter((elt) => !elt.dataset.setup) )[0]; root.dataset.setup = 1; const msg = document.createElement(\"span\"); msg.style = \"color: #cccccc; font-family: monospace;\"; msg.textContent = \"(Loading...)\"; root.state.loadingMsg = msg; root.shadowRoot.appendChild(msg); root.state.chain = new Promise((resolve, reject) => { const observer = new IntersectionObserver((entries) => { for (const entry of entries) { if (entry.isIntersecting) { resolve(); observer.disconnect(); return; } } }, {rootMargin: \"1000px\"}); window.setTimeout(() => { observer.observe(root); }, 0); }); root.state.deferring = false; const _insertNode = (node) => { for (let oldScript of node.querySelectorAll(\"script\")) { let newScript = document.createElement(\"script\"); newScript.type = oldScript.type; newScript.textContent = oldScript.textContent; oldScript.parentNode.replaceChild(newScript, oldScript); } if (root.state.loadingMsg) { root.state.loadingMsg.remove(); root.state.loadingMsg = null; } root.shadowRoot.appendChild(node); }; root.defns.insertContent = ((contentNode, compressed) => { if (compressed) { root.state.deferring = true; } if (root.state.deferring) { root.state.chain = (async () => { await root.state.chain; if (compressed) { const encoded = contentNode.textContent; const blob = new Blob([ Uint8Array.from(atob(encoded), (m) => m.codePointAt(0)) ]); const reader = blob.stream().pipeThrough( new DecompressionStream(\"deflate\") ).pipeThrough( new TextDecoderStream(\"utf-8\") ).getReader(); const parts = []; while (true) { const step = await reader.read(); if (step.done) { break; } parts.push(step.value); } const tpl = document.createElement('template'); tpl.innerHTML = parts.join(\"\"); _insertNode(tpl.content); } else { _insertNode(contentNode.content); } })(); } else { _insertNode(contentNode.content); } }); </script></treescope-run-here><div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtWgtX2zgW/itqenZIFmKc94PCWSfk1RZaCC1td+ZkZFu2RRzb2EpCmMN/3yvZeThxUtjSdmdn4BwC0tXVfeu7Eq8CNrPJicR8QgLN9cjAd12G/kCeG1BGXaeOfGJjRifkCBmuw7IGHlF7Vkcj13EDD2swPrUoI1nxRx15PozYNGBZwTrLZh6MOq4DwyrWhqbvjh09q7m269fDpUco+ku1gQD4UZ1ZdWRQBmQOIw47QiPqZKPxnCz/A3i5d9mA3lPHhHWurxM/C0NHyMO6DoNZmxisjvKaxaVxSNYi1LRgJCeV+H4OwxSUW/CPfslOaEBValMGKuIxcxe0WeownzoB1fi2JJyN9Hp4dRja8dXCjll/7MCePowFmk89hrghjvew59lUw9y0h67GCDeTT/Bo7ySdzhyfgOVhv4AhnRhOgI4Rs2ggmYRdglvOXZ2kM5LlBkwS86AaYWjgEYerrGicK1/079+SZrrY0W0C087Yto/CHSQQs++6Doymp64/zKBVGdxrGOJTsWFGNT7oEd9w/RF2NCI57jSdEYEAG6Q3ZlA2XPQKFfIZ4EMNlF6TWrKJYzILHR8jmZPsFN0nbOw7YHdE7IAsBbPGDpdsnXVgUYNx+QQB/+UBvrfskIbwc3R3KvnkdkwCpjh0JNzV9vGIpEObZDiPo42NvHFghWY8StBxvsVxqMYOLR8vA5cidCRzTdMO03cgUgyi1eO8+Aix2QEiEwjwyJNcOvG3NCQzbvSUn+ICRcSSZuMgeAtZHPFNpxY8ByMIw9R884cM2BPCX8T4yavDpATQ6QQJhsepeJ1JIYZV0JTcHafkFKSuzzZJXAdEBGM4MLUrGZItkOZr5rqnIBnDeqe6jLkjXhjqjsvSkuHaOlZhtQNs6xYO0ic2Vol9Ep8ZhHuINXXNItqQ6JkM+ic33bzwMNerI1nKlchos/TwsfusUBn+Pkossw+SKIcDrKo+mYjoFtXxZbmax7K8JNDc0QjUWqHA4ou7Zo0EhyJb7oT4mQT6kJzoAy7wkmFUkDfr9gT76WxWtV1tGA5ljuZVWOif8+5Q4NpU30UZOuFrxA/cvcTnwnkgDrHAIcRfVbrCv6PjSfi3jijDUGb54pj/thxsYK4kNwO1TgPYdDY/wNYJ0QkSYVKvqwRK3orlXmriK9nD4dmUzfHDKTrUwK2LvagjwkaYYceewpmbO+vYHwYEm5BGzubqZwr3hQx8afKiOX1MQnHI1tHer/mSqu39TPHii7YKWf4BQnI/8o3HfsAd6LkANYifsC8Nnm9bkQpio6wojsG2GH+eXZfqMXLHNneRaDAwqB+wgeuERWgztXalkpQv8WxKdBX6ZvFDj6+LyLUaYd8EcBiKIRL64Rt3g1LszdQxlEYnsQAtp5OCNoVSa1RgSIDmycS/klxRT60h+70zDFFBsY36s5Hq2gF6N2ZcXx01w5Xw6c0gMbJTog4BpYeVdwQHryXwOHYYLKc4IPoC278kMv8+2gzzcLXA1LJU4ydkXMswPxK0SC53y5XSFAcDDc4BMOxiPTZY7PSY1+lde66tiW+5avroANMxw1nsgGMFasusDvNNOBT1sTOPZsEW5QJEwGIADrLumD1NlYUE4BhK9BdxScSW6AUdea7PsLPBW/XdIXEGfGRZjL5u3ZVlK/acu/lBsjRrIJrCAc+OFYQUJosshRm7QudznLRC6Ie4aUHJcROoqg80aDR0nziR8vFWDqSMEz4bxFskY3Q0x5Jfw7aWhn4Tmh0ANAIfSwHDfP1C3u8mSYSoQkl0wFUQJhxWrbjjdoxtB9qHAXToBr0DJrHEq4rEg14Jc6Q1xb4DHhjMj4q5dw0Da7lCAqEHzcaG40Q9jIwUDWUjZy4hZV2AZOxnTR/rFNyWRrlCSSfmAXIhSUyCZBCvrFkHYdJAl8BLkBhCkZk3ZNmo1c9zEKCNkp9dBOe6r7fF5jrdAnqHPoTsjpP88OiVpaJmZR4j648J57A1EFIlNgixxkvoAT6xLfBWOgfqrC6J+AZARTK8gY5ssFyaffza+b7zNkasXLCcj2aj4R9ozWUjJiyY5Xk1DiIFc2QE2m1XLAFDPRnBfC3E/4j1iqF/ebV6RK+4m/hBAo7fpatfOvlbGvupBQcm7+uTkB0H2mC8RJqwJVyhAlm3MQo5gEls7AEC+3oX+3QHb99hKWhIRO4AquhbaZ5DjqQtYqZYXoclmSJ+3yTF77nQDg7rqjqxm7dkwsdstpXPM1428/s/9ELxfTyTDN8dpXVXG/ObIomDgkCaYHtMINAyUuCOSFpABX5zyD+lsFHgt4aPbBVSe3BwZhb3tIFFCOOXuWSKmv1+n2vT52P8alZMSj4Rdz39maOlf/9X1J5oZA5ant6qrF4OOfxy2o7GplEOF/nNWuBrdTT27TTHzXU+fzh1DSN/pAIyLxcPdLnWOTOVhiK+eheK4orfGpdT+NltK0pL2fXVGCmKOXTf6L1Wozn9rChXn5uvlbNeo6m0zbte963FgsYZJWahffop/7ZX/jzpe2P6/qx0lXv9qXf58WxyfXbP3s/a7eb+tTm8oo1T2aKnF+PXLb1zI3fVQ2PS073bN2Xr9prSi/GZ07G6xgemfCg3zv2i0u45w1ZZ+zAeO/uXpVstGE4nRts+vL0zW27VVF9PO9VcVzl0lMvSW99/nbvcN+/lS11WXhs587zSnHZu8qbszsaXlcqolStPu59q70zTI1fDWZH01PuSpvrvOgwr5kXvfHqKg1lwMe71Pl232lPl/YXX+6x/ODzcNytXlU8FJhtv3t8qkxLwfKucV5SzqTIy7y/7++MvfdL6dJc3ytr9efGyOyuNG8qb+8aN1/YKtHvRbMlfxu+L/YpjNN62uu2zkUL3q5NW3nJyVmVf/Tj9dDPt+pPTzoemc2O0Wibbf6d9se1KqdZ8PW1UrVrx7KzTL3S+KOaoV7ppXNTYVYd0a61Go9cpnJrFy8PP2kxVOuDTj28OlYsOVshZ01a696135hdmlhvvzXfveqeNIb0okXbjU7PR1qjsWb7rORAb3pfWae4+N+wbTYNZszdOV8ftoGvI56NO67zc0JXbjx89zIL+l5GuY1rLG/e14gd6c1v2Rn75nfu52ad+ZzR53Sn0r/uFdiuvNS6Mq/2u7XqdYjuYlrB5W67SL6R/bnvXTqPbI/qZT8bXt53mKHfd9of9/l0pX76+DqYKSJRB4mGGpfdEWO/xI/N3+LHIfqy7HvQOy5QUz0mSJO2gOAhz9jfgtfuC3hLvG6JhDHtZ4A3h4WgoHbaU8dcnSMErl6cvkEUtJx8LoDxwFryH5o0nnmLKkIMn1MTM9SXg7Kku9nVp6lNGrsgdSy95cUQR8lo+ccARn06tNNj8cQN2uaIjAp14ev76tbHOJyPoljeWPhygvCzLAklB8QXUmhZXRcn7rnTRqaVw/JJsXsH4e1AKvURtTG0obMxFnPiFqGyANh3o66AaU7AZwTq/BNhftV30UPOVJxp+oYBEeTxOxUBVHblDW7MAUleqJY5+5ALKVSpSvlYplAv5ai5fQocAgEDdJDzJL7ZTEfvoCSh+mb3epwFxCABeUccbRwdZShz5qnuXSmQSoQOYDJEB6CgWx/eNvXBEMBHF7yjWJI0f/amTX2zG7QgUu+lik2ttN8xGHMKPa3H8BOnFaKT9Bv85ykklaDV/5kmdQJi8xz4GoFqTUbpQllEjE99wdfXeHEXuxYf50MCe2PHh1EaHG4f1qR3UG6QqxcHx9rj4ljBIdv/3caZh4zvJce4AMAEMAUlsqkpxgwt/pJMcu7e8TNtDrtPkJeR474k1VDxwZvbQ4qbvOCVx46aW6byYAsAk5vjV7folJcyJegHnkQW/R4qfbA+eJ4ZlAeIxl3/WeNy4Dkztmk+dCFB7/AhvU/04tXh0rOIyHMx6TitX9KKq13BerWK5WMxVy5WCruXWNk16rIyiLVnTG4ggAcRRkkUW3aAgnJMYtotZIZ/OyQeokPnFXIbyV9wV6y0S1cWVsqFW9ZJay1eLWNWrOVzLF7SiQYrVfF5Xd1Xz58napBL4P2XVr+fCwVPzZVfV+a5lIuz0dhWLkOJ5SsZ6WYC68JEGUFTpveCILKpDzgKYQcsOmLfG/1Vc/1+Y7/EfWwJ8lwO2HYKr/4CSOgH4/whnbuZLZvcBvEF/sEm/uGOCYHliTv0JT9ftQb4bTA2J7xD7bzj1nRwemneby8PZnwGpAN2n88U/IaTKk5JcrIEIVcBRtVIRF4tVGcuGhtVisWxoPxtS5Z8dVZWIXsGVmlrQauUirhSr5Rqgq3y5WNTBALr2F0FV+b8GsApLwm5ssErzN7j6iSb8CwGs58ysP+lp+2d08/ZrzOW/saO1/6JKPT04Hg9L5pQ6nZz8B7KrUUQ=</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_5b83677f05674a588929eaebe4d0dfc0\")) .filter((elt) => !elt.dataset['step0']) )[0]; root.dataset['step0'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtfe162zbS6H9fBet+SKolmaS+Ldv7OE7SZpuk3qTZ3axfPw5FQjZtiVRJypbi9f+z93HOBZxbOJeyV3JmAJAEQFCSnXR337dxU5skZgaDwQCYGXzte/6NESfLCTnY9vx4NnGWe0YQBmTb8L2D7XEYnXtkTKKIeOdOrzse9b3OaGD3287I61vOwG657TFp923bG20f7sczJ4DfSO+w6USRs7zxP567YZA4fkAi4864vfQT0gA4l2BG0dSZDI17Qwfc9INxCChj+NIYO1N/ArxNwyCk2EPDDSdhtGd87dCfoTF1ogs/aIzCJAmne4bZtDtkOpRznEVkdXZ+MJsnp8lyBgKJnOCCbJ8BCzckSnzXmTSciX8RABe+502A0tifJAR4uABqMaSTqlUzQsjKT5ZVs9mpPTizvcvwhgqqSPph9IL5dEQiIBiESXVvHLrzuAZkR2HkkagROZ4/j/eM1mzxaSTZM2Uayad10qM/Q57dnmHNFkYcTnwvT1qRazMGSBLFqr6sqj3KQuLPAEdS5KExC2M/8UOoNmcEPMwT+DZy3OuLKJwHXoOzTDPSMTyaACxQcTzPDy6YXrmXSNYPoIYa5IYESZxmdut7yeUe1F7SQOYgaWggZ+NJeLtn3PixP0LFKRbrY8MPPLKAnE3TXF3KUbjYsJThohFfOh5mbdL/sFi0QHX+wYYPvOj6AmV8DVaw5U5899pzEuchNTYJHZTo+ZTEsXNBBO1JW/T9/i7rS/aTiJDYDWekEc2DxiWJ4FvsRv4sMahuVpzZDHhwUAK7oZuQpBEDjjOtHG7hD2QbJ0bKhXFgVKs14+DQuNsyDPg3ngcuohoeiUnkQyv/SN6BNPpVbDQAYBgRSeZRYNCvR0inOY7CadVJwhEA1Y3qlBKcNt3QIycoyqOkatZqQ8C+3yrP5jmIIWnZeUaM1dEyITHw+aj8UiJjpI1UAnJr8IworSol3xzNx9C3cxReQIazjusXwb+EZ5rNgzhmLE9IYhyjMk2d2ZsfnjwFzRyqpbkgyTEoox/Mw3lMgas3zmRO6kwNARPR0hIixZETk3PaGupGOB7HJGF8+GODoRr7B4aZYhgCPBTHHPKvDDP/cm+QSUwEIocHhlVCROSsOSHBRXJpNAy7QNpqysRTYkzEbhJnFFmW3xtVPWmrNtTx8cpJLpsgd5BZRqxW4CLP51vD4vwINR0pBTrNszg7Nc+QKQtYYORqxg4nb5QhGTuGxRHF2mGZXazKzHpsZpY+s9GqzOzHZmarmXH9P43qxkXdGJ3pG+0yANPJPYrc2A8u3xCgXuX5XZMl7fP/zNV+4s9eOaDbkXP7yg/YX3znJH5wZqlaZtTjxIHx7C3aJx7PogoYiTPPNRg1+ys/fu4HMDBUadLf/85UCIaq6qJm7CKCsW9YpNHO8bICLlLNUrQ5A6C0YrDMkNj3lNj3GQz+UIBJeFEt5rrDsX+NEqgU/jYLb6sLBlA37Fot0+17QYuzti/J0TiQ+gBMZ/JUEhTxo8Yw+ZcWVoVnZHWcqZBZqaWEtBeltZxCTZ1FNa13zlBtWFLS/Qwi4/Lf9Ic1O6gtbDPGwZZS82Qxq2YqIMsAmmGu3qA7GVhatKzqszxAVGlOu6why2LZVQTHBg8QLPS1gKtrMlnTw9YxFFQnxVJUJyMGLcYyZRXIulds6jOCvSxqO4y7qxrrLqeZ0uItKyWiHU0+ubZypU8lmb/qVFUlhHL1SBBO/QBsjCjTYT+oCiqgK7bS9XERUBaE3q6+hkraUBBdqjaJKaXuZIa1FciFv1G1CeTyIbYwCiThWzSk3yYRWNxsrJdtt2wQSYcuaYj5EF2Mqt/cRffGN3cX+Gt0X/ugHW7Qvo+cGMyqi8flKECgOxKAV7MEiKZt2V1onxF00c2e1bHh+QKfzZ6Nz6O8k8rRDg3L7uey54WpUAeoolXpFIQ6MRVJoEpBwT+fnDgJOEwB2EBQH/BvWQc3kX7KBkqo5CqqqU8NPvizn4JwIwu+7ezUFAstCm8BngOe+mepgmTkrhi5KyAHsBmpK5EUH3zC29OrM/ErZJIsmsj/G+ImVbQuroB3+OPXDasu2Hy5Rt4XVIsx6vkXfkKN55PInzoRVtUpha18PaY/lTo8WuNeb9Smj+Nxb2wS+mi7jmm79NHr2j27Tx8H7W5v5NHHvtvptkeVOidIWr2ea9OUkTvybPZo9UbEHVcAhopJ5estgS+ezFlvjP9RbIe4PdLnnI1GPc5D3xv3Hf510B906aPbGZlehz22B+6gnXE27o26HmPHG3mjPmN/QDyHdDLOtjLuXDKZvAUvCljqDVmC4rSAZzL2Lwo+iwc9zs8BOQb8tIejugeVTN2WusFdGD+Gns33cl+GEaxnHXqqENw4o9CCFnIFoU0YuKyE1xP3strpfItBg1pluKVRJMgKGmKPcsMe8F9tuJpmz1RoFhoWp4tGcEabvZxmenpq1o3831ldSrDoV6uY8FkwzmoFP04WehMjFDAMQJHdzOGsqE3e5Wa/ZL0USKVQ0mjjx6+d18x5rIkNvSBxseN7aPVpyOGgYHc6VB7wtyZQflQlgogtQcTWmTj+llZLVjFWoSpXY5XldVZTTPDcNYc6fAH1AI7MUhI0q0GsKpCLJrAA/njdKNRdUaS6YfpzVdfKkfnTa46J2Hp49TwuSZvX6pprPKrqzN9F1eUS1VVCWfNZ3R5XN1VrTaNTq4gPXB46vnwM413eULZ3UrhDjIIpds9j2+lDq/uhFf7oKn90pa9sYOsSLaH9WQ/CLE88E63OvPVm1blvmA+uTvP3V53mKrmbj67OB5PVVKdYdbnZklZwTa7bzMd/uO2hzbBAQkrDH1kFNNqUciqolFR194/WF8XLWmdIGpWZA85CwqPDlUL47ZMsSrlVpcmntK89E81ODgL8zIENGFqJ98X6/PzWpzoQlsRNdDLKwjZ1GrepC4Gbh1fBxspJfe9RWKqX6MSeOFESP1k+RdDMMadyER0ulF/vrPgFqs2uGy19Cgi2sxaiyyHa+LcFCrQWssMhKUb7MRgt/Nv5NEz420V9yb1PMcbsBzfox4NAxw7UkthYmRX8rWEZXynxyNx6yrCTaE7WaGFALpzEvyHZFOJ+PsOZwkydCzC35540BaHabEK8Bmd9M5wm9se0K64149nET6qVimLqMaR0snK/oFg8RWc0IOj5DGEhTxXvVCJ8JvbxKOVZRFdInEdkRpwkPg/HOFs9n0ykcVwT+JPIDo2dHV8d83gLjxPgpm7Evkc4D5xLxrIQECzIEABf02U6XDoAW5OBM8nRrrtYGJWnYreiCa4xkSlsGZuJqmS8XpUvDzY+JFcKqmSrtQ9oX1iwDvIRhFdOPooUakt8LrUQFGMobU+11VZLcRQtmEplg5/Qr58qxTvV2365k6dN0BiEZyXF1VSwWH5txyOaRWKwiwF/sTQ+p6WhXSuysRg2F2/JVAqaEWwq7ijwXgSe75K4qkayffYdH2LQE4qqLEBKO/RT3iXQTjidSwIkA3oCHXaqWZByWkE7pnJGLRkYdUjkTCqistE8mrN5fJkiUEYr2thTkaRqsaesZ4tmeDFPOaqz8OPK2Zk88KXABwaHiq/92TnthyrKVI/A7odv7jTg93vyZxJ48PGD3hvnGe8/NF+Gp6O6Dsdo6PLClUk4X7Wa703qg1WeOmlWmBApskcRH1CiSnXqxNfEM8J5Uqs8is3zSRhez2cFbtP5G+O774yvOK5/EYQROoi0t1xRO+V8FYvDVDWej+IEbDTadjMVZLyd07nqypniLqacyqhlnmOBw3lwHYS3gcReidUg4ImZlY1Lm8gem6BO9Gjd6ZttE1H0bfZwfdPJaD6mBRT16kHVxrJUWN+01tbW2eoWUlJf92WOiJjd/h8OK4pTEU5Ik0RRGFUr7xgvYt9f4eOIdmUXXwXAMrgK/SD1PaQFpkdQyW9nxC3M0p6D8ecs3wWJP/kzW/Bd9QjG/ujy5LrhULBUdsL4RZeH+zDqLH8exSS6oYt3+DpYEsWE4qVJ1SrYuJFP4mwhc1pf/Pupedb0BcQ3mD0ooFkcffhS7FdOdI3L7g8Mgd/mr3MSLd+CPesmYXQ0mVQr6tJtUfSscFVxWiJ1hcgEG4ySmaxCANKMyDS8IdWaTpWLEmp6fgyFCND2UCuzbtzdZ2uFoRhxchSA44AMPo+cKREWgZcQD9mDWH+pKaNf2T2a+xPviK8zf+5fzCOl8l0aLUlLvU5VZAbPN6Uusyiq5gP5S42ncciC2BjowTdxtStdVf/EiUm3nQMJHwuwT1moSAKl30RIOm69gqFJpawkiDi4AwGDBh5dhs7hhY8i7CJtvgJo/k2EXGogl1rIeAJjgKcBVxJEHDmKlqO4ykoQwcr1F9ATnpAIl4HkCNJnRZJz8pza2Dj/8EIwgCWplgENt/J+Cpsxra98uZVYpTBajtk2A2GkpACsT1E3PQhaUlxFoRD2g7Vk2a6EUqIpZnKJvgn2q8/Y+DAP4vlsFkYJmEEeHflrxeXqVO/O0ViSM2X7RBStrIlCyyQXTtDJ5Uvh8YM7jyLoreWPMZlRJ8YUvRgllJSt1c1VtpkGH5bqp5qyxgwzYOMmzz+P/Wf8IKs76Xu2sp/yL328V8splzgJE2dyHE5ipdzh5C+4UYqW0zrLE1hxQKoFN04RQKHc2nV0MQAATg4shM54HBNX5EJKE6wnsKnoIzWxUrCcIyay3DQRy5Y/fw8kd6iIIS9c6R9I2yiykjNqGZ6kcJnk3oS3iuRAc38k/sVlUhDdclPRLR8iuuUniG65WnS8cPnzGtHlRRdkh4i1ElXkQ46LcbW32AP/TK1bzPTuvlw+Sme9gZAUDEFSxcxPqZwCsD7OhD1AWyrXrgOeWyxO3VW90J1PoeE13Yg4CXk2IfhWrTDQSraNir426UZEXPudq+aOYeP2iHT1oQR+SSWbwdP60MMLQj2muDhPShaJxCunymdRIbVasT3KoqaT4Cuu6d6Un/jCbClOnMfswfzE4Q2rUDaNhFmQdOR8ky+AF0POq6autOvztCicZ9l65ev8IcMC/FTYDiAs7i+sDuer+9WV4YX5ZbYmhu4AORrFuhyzRHlhQoboLFYg0sQCw/qKOmDyLc5lFKqBL2inS45RILjkHhdvq8LRz99zA5IugH+VklyxvSn/0fFdx1zruQjruVDqhtm0OrXhpsWReMKPO7iBbzffl7c2GM7ny3CTkEZ1/EBThQ/RM7b5aCpu1FlXKLPZ2aBGVtRwA4tDaxg5ZW9rgjBbmt0dYLy/UHdQCmFd2gGDT6n0wNJ+i4zEzoHoIYgd8ferO+qhyiF3IfLxRDBHUANZ32QqayqWOfxyNTzvxvKh38Bt1RMCvXOEm5buhPkTbhUw2OxrKbj0KDgIgE8N59NMXmdFpUsDjQeSNcw7AOMPdP7G2DO++ipPLqGnWckuRhCUseUBK9y3CvN2ko5KOqh/ZMaUd+MELjkO50Ei6t5jjSpuEaW6BfbNjlC5mYGDX3M7h4KlBpEMW6rVksEmq2/eI4h8HB4IhhpOFxfmYZVXWTbFYsjFVGQn895Yyzs0TYk7lRbKDKhUS4RWW41fMvGsvI7A3rounzq+V7pbRXcOMnFo1yMouRXyEscNY6u4akPqBxblavsYJ4r6PytUFruonYPcpynV11JtXazWVhTnQtbVxSpdlV4Wa/R0sUpLS3V0odXRRbmOoZBQQ/VSqq1C1qrnxuoima8LVSkXq5RyqzSHssFa/qPpwpsTPyB/4U6JNVwBGCdReE1KJvPLKB87MwSOf507EVkL/ceQmlqVKU7wVn7TIXZr5cDGC6tbc8Im0rOtYg00yeqa3m8dTJb6vaiD1pnxhz+gmYprE1ZgCP1qGUrZkKobR63ScdT6Mo5+GUdXjKOHn28c3dps8LRKBk/ry+D5ux48Dz/D4El/i8Gw/EQVkkhximpAbtNneb2SkIADui7IUeNT2SkLZZG2nNIDo2S45GR9ZEyM3BUmadWIYCYJtjThBKe18LyDZdk2B+onCycS0eADDcCkAxme8mENNTvsRaxlCZYsEFznjMfXID5/xK+HYkPOY/z5OUpnALkld7eH0hAlRLcFLM1pOGIsc7UVlm6ci+eTRIh5f7bISjo5zIjk8X9ooBuFUwzOmxwSp/SGDwzfSEj3chQFtR3T+XlThbMUfDl+k0sf+/lD5iw1GsWCr5pcymAos8b6bp4H9UJvPpnHMjw9/SnDoW8inlq+/PXblN6wsN01jDzlGC8RcZfzXVimn59DhgR2NHM+aUPBdN4+GNJhPlkkd5catRb7399eRaDliSKEV62KLB+vIsuNVGStvarqiIiwUkmKJXyckkiIvx8lSc9ZU+OjdUMb5qxzdvKg5dnKCcr8SEyM5wvnBOVLsDB0GXjH4HB6pdOAnn9TqclnKfoBpSrMy6kzKizjh9OnlCU0NuOXTvDlpOkBntlsZEU4GLUy1IJmM5EbwOLRwM/pycDUm0/PBi6BTiIniHGx+c+Rf8ECAEk4gz5gXEYfBq6TKJyRKFlWK/7UuSCNiKD2+8EFHvFC19yAkLxKbQMCjUZ6AGnjYxhOkYC1ISKeNtagByHHYJlQzPZsUUnFzapDFvUH15m41Rsnqir5otX8zZ04TXw/W6S7AkVKWU1sRoqBl9BKj6elGwlQZmAeVFJlYfgbV5IILsmJnpOMwtGxWilOPqtC52f8voS8slJ/c1fo+O/pAY7NDpliPwvlTQtcQu+XcCaQW2xErrCKYTJ56YzIRFzcwWeV6PcTZesDH1tw/n0K7YLC0JAanTqf4Ks0f06/pBqEh1+9xYaE0p8J8TQRip2U/ISeto1wZtOCEghnJmuxUIfZ0UdCV1ShosFj2yw6LVs1m13sXIuVWEONYwlyk6il21Huf/s1Q7Q8yq7q1b2mThJii0iPbNaKbCKqY3kj5F6L0P7U+uI1xQhZpvkt6Ns3dz7qH1U/PR7vT4TSruNENV3zCKWGOYyCHrNOPnUiaIKOFYQ9whPYUWZCj2BoFZ1reebRFoYsOT1rYWx1kZh4/5uvpfoXq5TUa0ZhAuSx22wMTI9cVLS0Nf0y16gIu31tNpE0fmyscTpNbRiP1/wHa3DRsF6pwhrllGWygXZuyRY0qOrLR6tEjr62dQmgG6hQAXqU9/7lQKkWrIKBvj+dG6pgfZdDSqqbKWUjpFqJSixppaLKYuvPKT+2i9BQEAZjBnNHgfaEctRxpxV+orTva2UrAGng5Um4IHG5Bf9JPkKeQdOdOHH80o+TJlgsYOkG4xBFya9hyCynh4SHOBg75XDlMlHuBbFVxW85U4riKzMJAuubFBuvJakIpzgreXEa1Q//FXxzl7eR+9MPygIevACjwFpZpvTmDKFVMmQeRTUq7DKNipKaLsUpSokDsDVi5cnOQha3pQDECZmJcxt6UTBxMpR1UqucGRVFTkxvHiknhpzJid6LUlESy8XEAcrElCaXiokDqGLKPhc9LLRX26Z7iaerxHyAyEeIVfJlRIuyczFwCinH0N/gRSX0ugpyk0gGe0EWANCEol6QhH3KQx+KbpUDbskz1eXLlvkREyolqXkLcaGS+L9a5dD3PMODG7AjItBdprpR1wlFbVePQS52uhjFmQCFz9jlcop5/5N+adKFdJXM/piyDXGVwt51POobulrcyI+H2jQ7dE+9hYfQiDMEcgl+nvFhXOI/y3sT1kNKQuw4ZeKZRmXnhov7ZdCcKnYEG3YDObK2IxCSWUuvWJoU2sgrtiaJN++KEyyLiWmpNBuqRBIkOUqSyB/BYF6t0Bqui1WrBOLGId609fkCfJyiduRWQFY79SlU2pv/LcTLsxoqEZ6cC0CfXtmh/jzdmTwOa8U4RnZNU2kDKxphGwkkJcwk8ho6Kixn+rUyVHlQBvLPxwN2RkUe8DwsBehh4bqi+6EQ2jhaV6TELbDsEqfPGIdOaZaZmBmAHJ4TjllfvKL3usHAofR0GADChVM01o/1BePQE7zSyw8ujic+MPNG2hQsTPkE4Ba9ScX1zV1KibsqjYw0jb2AoD5oZ4TSGIBg85esBMlM3DRMJ+DIUzwcpEl9BMWByrjWwIuxVr2LhNIWN2pjKaidy2WokFojTDaFk+Gn2neYi459UQ7fWseuYufL3LNARZUGKWoGTcLge7VhNjtof1lNm0yVQyI+pYSGrCc8fltQllwIXFtw45R7KZ6zp1mL8mjSphw4UVeVSHo+nk/Swqe0V2ly7t+u02NONEPA5Txr5SlwI++aGid7+c0dKRQtcj3NjL5JF3Ik4UyHBZ9zJHiRcKhc9/LdsxkWTcjx6KuEmd4qWURlKTkuexeQ7zX3s6yadhD0gIZ+G4YkErE7WjXhUAh3sQsxaQQuywCnVgT68CrFt0roC12nLEGx64z4bNA6YtkMgkiNB3EEcuxLTu9e2ZWOAxNQQ/OFLhmK83UtTGvnMxhmCKY/j8LpW26Z6rcRSqsD8LQT72TBjs/g8+H0a/XWD7zwtumRG/AwaK4USDqcvwQGb7uSFveo2VjGrpgVvq7LTr7/Z5YBHHlX8ziZshigklEp1fIbd9w4Pln8Qt0+eklJFLP9/FXVhpbkIGIpZZcJWmpxEFtC/14pRM34VrhS46BwwMmnTenqj8B8EE1nnoSVkmqic5aoh2IJxUVo32trsqxRrZlwhjZGMxQOrhItLGkKjfF36yTupaC3aovhi1KI58sHDXO1miL6K0ytfqhGhAaU6S2l39yVKN69N8MuSAmLTLhbX8jfYHnzU2OUEIB7SZ3GeoYujEa6/kAcrZSCC3GL9IFlXAw8lOV6X3Aiy4MWciFXdG7ZAXITj80zC+TZaU7wOQVdXeYtqeZiNwonkyeSTXZHB11dDnidIeWgbozIpXPj40WuFTxiyQmSyr2ah+hI03xeBEn4Z5/cVu806EBzErrX8CUgDihRTvBe3vRelOc0nMdUM1CmagwtPfwI15mC5MQFpxjaYoX7a93IX95LfVqKqdkBnbmjrFnyy3pp4DcMiHCCpOLFlQGWHQad7YY8v1GOScwWH2Ul2yRTKuhy9thctGSlr3E36aJZuiG6ZA6tkEeiWC/rs6AbPtZngdUhT4jpT56U8i4X8Iojc2rDDRRBFbUQnADTCCC6ahpaTMRLJ8oyykw12eedjIAe+cc0SqBg8+86dHqkWN4coKWLfBwapvHdd5LIROAdBZi5XgLHsoOoSIsvfahI590qMNw7VjzHzaaO5WKlRVd0R+O8adnUcrBR1jtlWZcVVpLIfUlNvc9q6sfUOV9RVe+zqsqgxbr6UePNK7zRdru6rtL54k+srPefUlnF7uUBdfX+AXWVT45XSreAbDR8hcwiUEevDYeYjQaYtYzQON26MfRNutXgkSNpil8cTou7YhDg8YPjeT4qskyH6omq5zfFE2BXc6E/2SMLgMojT+UYE4i3h8FyRRvXnIesFqZw3InuZ7PBypBvzNJvKZJ3VK2TSclGrAKaei+A/oSQjXT19pKQiU5X0/7RmSSQbzGw5ZFJ4rxHF4Ph1prsS84VYvPj5Z+SsQOKo4ZE87uoC0Z7De9eyNMZbeXem5WWviiTIiCua1fY5a4BjBoI9xL4nmzosmfn04mIEqGhuhG6yNGh0V5XusaB0Vb1T8p1v4RdPM2mcJyNzK70uiPvYNRt9Sstx77RWFuQnXUFOSwriB88qCCN9QXRTGaKJDb3gYsPj5t6KTqVd2oklzZMl+K+F+N+GMI1dhUd/l7p86qKhCRoKY7Lgs1CZn9Vp3s+Z246XzavL81xtTMnjv0bsscucLmX5sR0qyhWVaA2gFE4sJYFqn4OyFN+bc/nO65WPl1IPGhUs/iVBS84UPpFA5eFi0VQ9nEoH89UPEJv80P0CsfoyYfgFY/NK6QzFhJ1tVj5wXjFc5HohRLs3gY++ONZE6AluCRiT7hoCTSHj7rKzthsY9EjVpGWbuGxXGdWKYPKd++sBJvSiCOr84rZ7HXItBRWsFj9AE/vaMiec2FeQb/VRA+MxvF4QjcDV26YSpfCrt5tpJFY6T4KFRJLlXnmFbvZqWwU310Xvi4r8zzBDGklzRbKlo2tVXvDsk1eUiNYqR8iSJlyiDAlC5C3lEnnX3BYKBxjjinPJgVffEONLyt1SlY0jvgntey2uHa5AKbd/GJgUB6bALia2XSmRdd5wUiUFjY9aeCeTD+U0l+5eLsAnQZXGp1VPOderGWtgpOWv7skEA/WKWYt6bulA5TdJXVAuFf1gY0Avx+VYOX9t2kFj/X821UiHfhlrShbspTVZfEMf77H945fFqDYRHvFT8yO0xynv6f7iND3w637GrXCkkufmgBvwjB5HXqkWmtehnEC3uY4iJtpACo9kREeh/u7YDn7s+RwfzeJCIldGAIa0TxoXJKIHO7jAneDrrA62B6HEw9v8zgPgPL24T6V0+E+nVgy0HI42HYviXsNZdjW4pwn4cXFBFF3KZJMnp72ce6MRuAFbx9+N0mGYnIlCBOaWDm8chZNKgQDmAcIiQz66EEOmILww+OrFpg8rVoJYeOf//hfZtM0/t//xd+n//zH/6H3/P7zH/8b/p4ZH0kU7rXMNZmmyekfLiWxpBSaLOCLR7xtOQ33iYMieucuKhYox8p0SMTVnzwtrdTzTB0VZOV2je3DN+kIzxSi2WymfGuVgWoKr2pQ/4nv0qawG7oJSRox4DjT7cPsyizWnKhOsjdUyWEWEhnTawJWquyQgTWBh7dhGKTXanASyYwHHbC3lS4TqVYSMp2huYJ0SBSBwMAwxA0nuQnLbkX449ufXzepA11Fgk2+L1ylx8peqUl9BVBDlOxeEUNuZ03dnRiM3fxCDewr1jTCtCxpNUpF2l5RL1dxCPpzt41rIrf3jO2nLOyWHawDXgMGC8FPNhy6AKDWNH5Ew3WXfcelXOw2BBpo3K4b28IdCEjx6D/8J2OZXveAHPOeABOUqxVoef707Mkm/yO6cNYDoN5tOxb8sTBDEx5a4Dtt53t94cvp3TZuc8BsABRSaddAXwHKsAAHv9I9DvAVep5taJ6U1Bk8L8tomUVaJjBi6olZJqWmbK9CkvhZPnKJFisVXH44ExIGwwHJYv+4PXUW2TOP7MI7cym3s3OwM5DsOOzsS3oXNHJxOuiyLQgtekmt1Wqz17ZFX7s9SOni//S1D489vG6XvQ76AGzivYDsLl8T31sD+GWxq3st/NDDyyc77IMNH2yzBR8G7JLcNm6GwBzsFoMwbdwOQX+xTLqIPsC7Ck2TfrBw00TPpF/xQ6+DLGNOfYrRwke7lb0jsUE7A8ccu1joLr2GtwXU2kjRss7OUImk7QIgoh5T3bKoO9OOtNb4TYPbKQ5VEWwQ2xgVWgl2KoPQE0xoewKlyRUvvzuMKlf6JVO2dbnA85psrEI2LTmX1tpMztaX9r8CA0YDGY4l3p/dix0074zhEYbewnBfsA0Uu2aVIfB34yldsrRnHJ+8M1JD47uLZKjPhTGABgC1dGFgYAGNPQMn4LYN30MDLDrPxou+0x3YY89yuz2vPfIGjj3qO2a7DXrZa3muxdmD30jvkF0VBNbEFD2Yc4nbPTCbqk3JvNu7dOLqIbN3mjrDj+LsURuReHjgQo0O5RMMY37t0B8YDqF0NPffxvY0ePab2KCaJK0sdLUK4BNnFmMiNWU1iqHQ+Fprzv53M3tXCqJck39jK5nKrXp6ajbxCgb8P7+u9r/jRxgn6HVSB7yCayVi/bTuokO8ntMbjFruoNt2eu1+d9BvO3a33fbag47n5t3F/0Av0V7bYjo2bTJ2P20zdoe3msHgDGVKW073i8P4xWH84jCmDtYvs9ny5KPzjFz13Re/hjvvTv56c3IVXT/7ZeeCvL3esV7/QE7i5ZNXT/vvXz7t77pP2jcn/t+Oxk93jn50+/3l6+urk/aLV8+e7l68mr7fta+sG6D3059/2e3f/nK9e/Pi3fFJ//iH9tXiYufiaDB/9sQ6sUZvnjzdefXxOtydT/qvT7rHT3582p/O7D/tPr/2nsB7+y9Pu6E9DXeOfno9OBn9cfnj1e67xUm8M3kRvzzZefKT+bHz06sLd7f/6/PXJ23/j++ueu+Pds2dzg8/Xp4kT25Pnrbbb+x3O3+bvr886c+Wb6923b/c/mk3fHb8p5Nw2Yp+6R71Xr/rv54nwcnuk5M3Vzuv3DdHu/2fX7ZP5sfm/KP96uXS3X0fHN2ctI5evb0avJ+/f7/z9KX1t5Od41+dq86rj91w57dyaD/J8bVTx9fixjn1gLsrPeBHea2lPrBdpGav8KfRyF/ji9t6XPszes8NGCn6nVan1bd77V63Y1odwZvWJK1zru2O3be7ttXqt7rd7mBgF7xtDdEvzvcX5/shzrct52J/hkzsLx7+ZzLZbdIx2wNgpQ9u/aDTdtrtvumYY9cZtdvd8X+Ayf4/y6n+l3sJ/9n+NfioAxhZuj3DQD910Ot3eoMBfW61bHsAHXPu0KJH2+72u91+u4sQHXAv+/0ehba6nY7ZMlvg6WaXaiD1lmVag3aXwnT7A1DuThufe+0uZNvpydRhsDGtdqdTz0c+ipkOlCr1bts02xaD6XTt7sDsU9w+DKUwqBoy9UGr3Wn1Oi0K0WpbbdOmmD2704XhxFCpd0zkweIcwPhrW5R3swVi6fcUyQw6fdPqc17gwRpw6gMb8GA0VKl3bBg+may7ltXq2Oy5P2hZvU5H4b3XaQ8G5oDmb1tWt231BvTZhlINegW5m13Iv2W2GHWg2WN10B5YvVbX7MjUO21aVkq9bdndvtmlmCB1GFitgmT6ZgtqpMW4MTtdkE6PUYeCDgYK9XYbdGAwoBBWH6TeZjK17dYAe71CrbZMULMel57Zh4pqM6n2TBvsJJk6sNhrt8wO09mBDVLqsxpDqfY6KvWB2Rm0bFZPPVDDTs/uMP1p2y27q8i93e9aZrfNdGAAWt7qULnDg9kd2IVa7YLGmH0m615/gFJiZe22LMxVpt7vQHMwB1waQL7HWlar1epaCP0pgaNPdvDZgV4vPHqasOCT9rrjUd/rjAZ2v+2MvL7lgKq57TFp923bG1GPY+MItgi8Nn4lAq8dOc+GPA4RulEYJtpIRBosQFlxIA6uhAe2m5k4zzF5u4B6PAnpSrr0veniB5oNvRhsmKdEBGwCl/zFTy6rEnpGdBw5F3ybvrI86Sl/fc4hsAwptLTEJSWM6ZoTPrLKVU74YFeOp1vwM8rwyhdGPVm+AOIpNt69M2SX1YfzyCVP6dlWJTL8Gk2fbWPHUND5yXGiWDJqzbEfxWnetGSAkKfmkZl7uSa0Ql4XiPn/U2qZ9w==</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_5b83677f05674a588929eaebe4d0dfc0\")) .filter((elt) => !elt.dataset['step1']) )[0]; root.dataset['step1'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "class Weights(nnx.Module):\n",
    "  def __init__(self, kernel: jax.Array, bias: jax.Array):\n",
    "    self.kernel, self.bias = nnx.Param(kernel), nnx.Param(bias)\n",
    "\n",
    "weights = Weights(\n",
    "  kernel=random.uniform(random.key(0), (10, 2, 3)),\n",
    "  bias=jnp.zeros((10, 3)),\n",
    ")\n",
    "x = jax.random.normal(random.key(1), (10, 2))\n",
    "\n",
    "def vector_dot(weights: Weights, x: jax.Array):\n",
    "  assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\n",
    "  assert x.ndim == 1, 'Batch dimensions not allowed'\n",
    "  return x @ weights.kernel + weights.bias\n",
    "\n",
    "y = nnx.vmap(vector_dot, in_axes=0, out_axes=1)(weights, x)\n",
    "\n",
    "print(f'{y.shape = }')\n",
    "nnx.display(weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2b222eb",
   "metadata": {},
   "source": [
    "Notice that `in_axes` interacts naturally with the `Weights` Module, treating it as if it were a pytree of `jax.Array`s. Prefix patterns are also allowed, so `in_axes=(0, 0)` would have also worked in this case.\n",
    "\n",
    "Objects are also allowed as outputs of Flax NNX transforms, which can be useful to transform initializers. For example,\n",
    "you can define a `create_weights` function to create an single `Weights` `nnx.Module`, and use `nnx.vmap` to create a stack of\n",
    "`Weights` with the same shapes as before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0b076a0f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<script> (()=>{ if (customElements.get('treescope-container') === undefined) { class TreescopeContainer extends HTMLElement { constructor() { super(); this.attachShadow({mode: \"open\"}); this.defns = {}; this.state = {}; } } customElements.define(\"treescope-container\", TreescopeContainer); } if (customElements.get('treescope-run-here') === undefined) { class RunHere extends HTMLElement { constructor() { super() } connectedCallback() { const run = child => { const fn = new Function(child.textContent); child.textContent = \"\"; fn.call(this); this.remove(); }; const child = this.querySelector(\"script\"); if (child) { run(child); } else { new MutationObserver(()=>{ run(this.querySelector(\"script\")); }).observe(this, {childList: true}); } } } customElements.define(\"treescope-run-here\", RunHere); } })(); </script> <treescope-container class=\"treescope_out_5aed62c96b39448cb9c45352442f166e\" style=\"display:block\"></treescope-container> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_5aed62c96b39448cb9c45352442f166e\")) .filter((elt) => !elt.dataset.setup) )[0]; root.dataset.setup = 1; const msg = document.createElement(\"span\"); msg.style = \"color: #cccccc; font-family: monospace;\"; msg.textContent = \"(Loading...)\"; root.state.loadingMsg = msg; root.shadowRoot.appendChild(msg); root.state.chain = new Promise((resolve, reject) => { const observer = new IntersectionObserver((entries) => { for (const entry of entries) { if (entry.isIntersecting) { resolve(); observer.disconnect(); return; } } }, {rootMargin: \"1000px\"}); window.setTimeout(() => { observer.observe(root); }, 0); }); root.state.deferring = false; const _insertNode = (node) => { for (let oldScript of node.querySelectorAll(\"script\")) { let newScript = document.createElement(\"script\"); newScript.type = oldScript.type; newScript.textContent = oldScript.textContent; oldScript.parentNode.replaceChild(newScript, oldScript); } if (root.state.loadingMsg) { root.state.loadingMsg.remove(); root.state.loadingMsg = null; } root.shadowRoot.appendChild(node); }; root.defns.insertContent = ((contentNode, compressed) => { if (compressed) { root.state.deferring = true; } if (root.state.deferring) { root.state.chain = (async () => { await root.state.chain; if (compressed) { const encoded = contentNode.textContent; const blob = new Blob([ Uint8Array.from(atob(encoded), (m) => m.codePointAt(0)) ]); const reader = blob.stream().pipeThrough( new DecompressionStream(\"deflate\") ).pipeThrough( new TextDecoderStream(\"utf-8\") ).getReader(); const parts = []; while (true) { const step = await reader.read(); if (step.done) { break; } parts.push(step.value); } const tpl = document.createElement('template'); tpl.innerHTML = parts.join(\"\"); _insertNode(tpl.content); } else { _insertNode(contentNode.content); } })(); } else { _insertNode(contentNode.content); } }); </script></treescope-run-here><div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtWgtX2zgW/iua9OyQLMTYzouEwlmH5tUWWggtbXfmZGRbttU4trGVhDCH/75XsvNw4qSwQ2d2dgbOISBdXd23vivxMmIzl5xKLCQkMvyADELfZ+hXFPgRZdT3GigkLmZ0Qo6R5XusaOERdWcNNPI9PwqwAeNThzJSFH80UBDCiEsjVhSsi2wWwKjnezCsY2Noh/7YM4uG7/phI156jJK/dBcIgB81mdNAFmVA5jHisWM0ol4xGVdk+R/Ay78rRvSeejas80OThEUYOkYBNk0YLLrEYg2kGg6XxiNFh1DbgRFFqvD9PIYpKLfgn/xSnNCI6tSlDFTEY+YvaIvUYyH1ImrwbUk8m+j18PIwtuPLhR2L4diDPUMYi4yQBgxxQ5zs4SBwqYG5aQ99gxFuppDg0d5pPl84OQXLw34RQyaxvAidIObQSLIJuwK3XPgmyRckx4+YJOZBNcLQICAeV1kzOFe+6N8/Z810sWe6BKa9sesexztIIGbf9z0YzU/9cFhAqzL4NzDEp1LDjBp8MCCh5Ycj7BlE8vxpviACATbIb8ygYrzoJSqpBeBDLZRfk1pyiWczB52cIJmT7BQ9JGwcemB3RNyILAVzxh6XbJ115FCLcfkEAf/lAb637JCH8PNMfyqF5HZMIqZ5dCTc1Q7xiORjmxQ4j+ONjYJx5MRmPM7Qcb7FSazGDi0fLwOXInYk823bjdN3IFIMojXgvPgIcdkBIhMI8MSTXDrxtzQkM270XJjjAiXEkuHiKHoLWZzwzecWPAcjCMPcfPOHAtgTwl/E+OnLw6wEMOkECYYnuXSdySGGddCU3J3k5Bykbsg2SXwPRARjeDC1KxmyLZDna+a65yAZ43qn+4z5I14YGp7P8pLluybWYbUHbBsOjvKnLtaJe5qeGcR7iDUNwyHGkJiFAvonN9288DA/aCBZUipktFl6+Nh9UagMfx9nltkHSZTDAdb1kExEdIvq+KJ6pGJZXhIY/mgEaq1QYPHFXbNGgmORHX9CwkIGfUxOzAEXeMkwKcibdXuCw3yxqLu+MYyHCsfzKiz0V4I7FPkuNXdRxk74FvEDdy8JuXABiEMccAgJV5Wu8e/keBL+bSDKMJRZvjjlvy0HG5gry81AbdIINp3ND7B1QnSKRJg0GjqBkrdiuReG+Mr2cHw2FRV+OCWHGrh1sRf1RNgIM+zYUzhzc2cTh8OIYBvSyNtc/UzhvpCBL81eNKdPSSgO2Qba+0mt6MbeHyleetFWIau/g5Dcj3zjcRhxBwY+QA0SZuxLo+fbVqSC2KgoimO0LcafZ9eleozcsc1dJBoNLBpGbOB7cRHaTK1dqSSpFZ5Nma5Cv1n82OPrInKtRji0ARzGYoiEfviNu0EpDmb6GEqjl1mAltNZQZtDuTUqMCRA82zin4hSNnNryH7vHENUUOyi/myk+26E3o0Z19dEZ/FK+AxmkBjFKdGHgNLjyjuCg9cReBx7DJZTHBFzge1fEJl/H2+GebxaYGpZqvMTMq1lnB8ZWmSXu+VKaYqjgQHnABh2sR5bLHV6zOv0rj3X1qS3XDV9coCZmOEi9sCxArUVVof5JhyKhtibR7Ngi5QIEbAYgIOiP2ZPU2UhATiGEvOHtCRiS/QDHQV+yLC3wVsP/SHxBnxkWYy+bd2VZSv2nLv5QXIMZyCawgHPjhWEFCeLLMUZu0IXcpy0QhjGuGlByXETqGoODGg0zJB4ifLpVg6kTBM+G8RbJGNyNKeS38CukYd+E5odADQCH0sRw3z9Qt7vJkmCqGJJTMBVECYcVq2443aMXQ/ahwF06Ba9AyapxDsSiQe9EuZIa4pDDzwwmB8Vc+9aFjaUUgZhAM3GhuNEPUyMlAwVE2cuIWVDgGQcFu0QmxTclkdKqWIS+wD5kCQ2QTKIVzWcgzhpoEvgJUgMocTMG7Js1OrnOQjQRskvLoJz3dfbYnOdbgG9Yx9CdqdJfvfolaWy4RQeI+vvE85xayCkymwQUo2X0AN84jrgrbwC6qwuSfhGQEUKvIFObLBcWnz82vm+8zZGrFywnI8Wk+Hf0ZrLRkxYsMjzahwlCipkBNptVywDQz0ZwXwrxH9N9Yqxf3m1ekSvuJv4QQKO36WrXzr5tzT2UwcOTN7XZyE7DrTBeJk0cUu4QgWybmMUcwCTuDgABPbtLvbpDt6+w1LQmIjcAVQxt9I8hxxZW6RMsbwOyzJF+r5JSt9zoR0c1lX1Ujdv2YSP2Wwrn2e8bOb3f+gHLQzxTLJCf5Q3fWPMb4okDgoiaYLdMYFAK0iRPyJ5ARX4zSH/lOJGgd8aPrJVyO3BwVlY3NNGDiGMX+aSKTrr9/tcmz4f41ezYlIKibjr6c88I//Lv5L2xCBz0PL0VmX1csjjl9NuMjZNcrjMb9ai0GigcejmOW5u8PnDqW9Z6rEOyLxaPjDleufc1pqa+OpdapovfmteTeFnt61pLW3XV3OkafbQf2P2Ws2z6WdNu/589lo77zXPtLZ91+u+dVjUPKfELrVffVLf9qqfJ/1gTN+fV66V1596Vx/PJzfn9+z9rN0+27+xh9e0+Up26KvL8euW2fkqd/VDa9Izg9s3Vef2htLL8bnXcbrWB6Z9qDYvwrLW7nnDVtX4MB57+1eVWyMaTidW2z28vbNb/pGtv552jpSuduhpV5W3Yfhaudq37+UrU9ZeW4p9UTubdr6qtuzPxle12qilVKfdT/V3th2Q6+GsTHr6fcXQw3cdhjX7sncxfYWjWXQ57vU+3bTaU+39ZdD7bH44PNy3a9e1TyUmW2/e32qTCvB8q13UtPOpNrLvr/r74y990vp0p1pV4/6ifNWdVcZN7c1982vQDkq0e3nWkr+M35f7Nc9qvm112+cjje4fTVqq4ylObV//OP30ddoNJ686H868r1arZbP9d8YX161V6mevp80jp14+P+/0S50vmj3qVb42L+vsukO69Vaz2euUXtnlq8PPxkzXOuDTj28OtcsO1sj5mat171vv7C/Mrjbf2+/e9V41h/SyQtrNT2fNtkHlwAn9wIPYCL60Xin3yrBvnVnMmb3xuiZuR11Lvhh1WhfVpqndfvwYYBb1v4xME9O6at3Xyx/o19tqMAqr7/zPZ30adkaT151S/6ZfardUo3lpXe93XT/olNvRtILt2+oR/UL6F25w4zW7PWKeh2R8c9s5Gyk37XDY799V1OrNTTTVQKICEg8zLL8nwnqPH5m/wI9F9mPTD6B3WKakeE6SJGkHxUGcsz8Dr90X9I543xANY9zLAm8ID89A+bilTL8+QQpe+zx9gSxpOflYBOWBs+A9NG888RRThjw8oTZmfigB50D3cWhK05Ayck3uWH7JiyOKmNfyiQOO+HxupcHmjxuwyzUdEejE8/PXr411IRlBt7yx9OEAqbIsCyQFxRdQa15cFWXvu9JF55bC8UuyeQXj70E59AK1MXWhsDEfceIfRGUDtOlBXwfVmILNCDb5JcD+qu2Sh5pvPNHwCwUkyuNJLgWqGsgfuoYDkLp2VOHoRy4hpVaT1HqtVC2pR4paQYcAgEDdLDzJL7ZzCfvkCSh9mb3epwFxDABeUi8YJwdZThz5un+Xy2SSoAOYjJEB6CgWp/dNvXAkMBGl7yjWJE0f/bnTH13G7QgUu+lSk2ttN8wmHOKPG3H8RPnFaKL9Bv85ysllaDV/5smdQpi8xyEGoFqXUb5UlVGzkN5wdfXeHEXupYf50MCduOnh3EaHm4b1uR3UG6Q6xdHJ9rj4LWGQ7f7v40zLxXeS590BYAIYApK4VJfSBhf+yGc5dm95mbaHfO+Ml5CTvSfWUPHAWdhDi5u+k5zEjZtbpvNiCgCTmONXt+uXlDAn6gWcRw78nih+uj14nhiWJYhHRX3WeNy4Dsztms+dClB78ghvU/Mkt3h0VGqGWa4SHZu1WrkqW3VFUVWlStTKkXGkkMraplmPlUm0ZWv6FSJIAHGUZZFFNygI5ySW62NWUvOKfIBKhR/tZSh/w12p3iJTXVUF1SxDISaoXTdAXaxXylYNqzLoY5V3VfPnydqsEvg/ZdVv58LBU/NlV9X5rmUi7vR2FYuY4nlKxnpZgLrwkUZQVOm94IgcakLOAphByw6Yt8b/VVz/X5jv8R9bAnyXA7Ydgqv/gJI7Bfj/CGdu5kth9wG8QX+wSb+4Y4JgeWJO/QlP1+1BvhtMDUnoEfdvOPWdHB6bd5vL49k/AlIBus+r5T8hpFKPVKN+JKtHulEv182yfqSaVegsK7IuW4SU/2hIpT47qjLNck3VLaIbWC4bunlUls2SqlSsSkVR67ryF0FV6l8DWMUlYTc2WKX5G1z9gSb8CwGs58ysP+lp+2d08/ZrzOW/saO1/6LKPT04Hg9L5pQmnZz+B+svUaU=</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_5aed62c96b39448cb9c45352442f166e\")) .filter((elt) => !elt.dataset['step0']) )[0]; root.dataset['step0'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtfet627ay6H8/BeteJFUXk9Tdsr2+NJfVtE3jJukly9ufQ5GQTFsiVZKy5Xj5/1nvcc4DnFfYj7KeZM8AIAmAoCQ76Vp778ZNbZKYGQwGA2BmcDvw/CsjTm5m5HDX8+PFzLnZN4IwILuG7x3uTsLozCMTEkXEO7PtHrEnrkU81+sM3cnQcsbdzqTv2Kbjkkln9+ggXjgB/EZ6Ry0nipybK//9mRsGieMHJDJujetzPyFNgHMJZhTNndnIuDN0wC0/mISAMoEvzYkz92fA2zwMQoo9MtxwFkb7xucO/RkZcyea+kFzHCZJON83zJbdJfORnOMiIuuz84PFMjlJbhYgkMgJpmT3FFi4IlHiu86s6cz8aQBc+J43A0oTf5YQ4GEK1GJIJ1WrZoSQlZ/cVM1Wt3bvzPbPwysqqCLp+9ELlvMxiYBgECbV/UnoLuMakB2HkUeiZuR4/jLeN9qL1YeRZM+UaSSf1kmf/ox4dvuGtVgZcTjzvTxpTa6tGCBJFKv6sq72KAuJvwAcSZFHxiKM/cQPodqcMfCwTODb2HEvp1G4DLwmZ5lmpGN4PANYoOJ4nh9MmV6550jWD6CGmuSKBEmcZnbte8n5PtRe0kTmIGlkIGeTWXi9b1z5sT9GxSkW633TDzyygpxN01xfynG42rKU4aoZnzseZm3S/7BYtEAN/sGGD7zo+gJlfA3XsOXOfPfScxLnPjU2Cx2U6NmcxLEzJYL2pC367mCP9SUHSURI7IYL0oyWQfOcRPAtdiN/kRhUNyvOYgE8OCiBvdBNSNKMAceZV4528AeyjRMj5cI4NKrVmnF4ZNzuGAb8mywDF1ENj8Qk8qGVvyc/gzQGVWw0AGAYEUmWUWDQr4+QTmsShfOqk4RjAGoY1TklOG+5oUeOUZSPkqpZq40A+26nPJtnIIakbecZMVbHNwmJgc8H5ZcSmSBtpBKQa4NnRGlVKfnWeDmBvp2j8AIynE1cPw/+JTzTbO7FMWN5RhLjMSrT3Fm8+us3T0AzR2pppiR5DMroB8twGVPg6pUzW5IGU0PARLS0hEhx7MTkjLaGhhFOJjFJGB/+xGCoxsGhYaYYhgAPxTFH/CvDzL/cGWQWE4HI0aFhlRAROWvNSDBNzo2mYRdIWy2ZeEqMidhN4owiy/Jro6onbdVGOj5eOMl5C+QOMsuI1Qpc5Pl8aVicH6GmI6VAJ3kWpyfmKTJlAQuMXM2oc/JGGZJRNyyOKNYOy2y6LjProZlZ+szG6zKzH5qZrWbG9f8kahjThjE+1TfamwBMJ/dR5MZ+cP6KAPUqz++S3NA+/xeu9jN/8cIB3Y6c6xd+wP7iOyfxV2eRqmVGPU4cGM9eo33i8SyqgJE4y1yDUbM/8+NnfgADQ5Um/f3vTIVgqKquasYeIhgHhkWanRwvK+Aq1SxFmzMASisGywyJfU2JfZ3B4A8FmIXTajHXOsf+PUqgUvjbIryurhhAw7BrtUy37wQtztq+JEfjUOoDMJ3JU0lQxI8aw+RfWlgVnpHVcaZCZqWWEtJelNZyCjV3VtW03jlDtVFJSQ8yiIzLf9Mf1uygtrDNGIc7Ss2T1aKaqYAsA2iGuXqD7mRgadGyqs/yAFGlOe2xhiyLZU8RHBs8QLDQ1wKurslkTQ9bx0hQnRRLUZ2MGLQYy5RVIOtesakvCPayqO0w7q5rrHucZkqLt6yUiHY0+eDaypU+lWT+qlNVlRDK1SNBOPcDsDGiTIf9oCqogK7YStfHRUBZEHq7xgYqaUNBdKnaJKaUupMZ1lYgF/5W1SaQy4fYwiiQhK/RkH6dRGBxs7Fett2yQSQduqQh5l00HVe/uI3ujC9up/hrfFd7px1u0L6PnBjMqunDchQg0B0JwKu5AYiWbdk9aJ8RdNGtvtW14XmKz2bfxudx3knlaEeGZQ9y2fPCVKgDVNGqdApCnZiKJFCloOCfz46dBBymAGwgqA/4d9MAN5F+ygZKqOQqqqlPDT74c5CCcCMLvtXrNcVCi8JrgOeAJ/5pqiAZuQtG7gLIAWxG6kIkxQef8Prk4lT8Cpkkqxby/4q4SRWtiwvgHf74DcNqCDZfrpF3BdVijHr+1E+o8Xwc+XMnwqo6obCVzyf0p9KAR2vS74879HEy6U9MQh9t1zFtlz56PbtvD+jjsNPrjz36OHC7vc640uAESbvfd22aMnbHns0erf6YuJMKwFAxqXy9JvDFkznrT/A/iu0Qt08GnLPxuM95GHiTgcO/DgfDHn10u2PT67LHztAddjLOJv1xz2PseGNvPGDsD4nnkG7G2U7GnUtms9fgRQFL/RFLUJwW8Ewm/rTgs3jQ47wMyGPAT3s4qntQydRtaRjchfFj6Nl8L/dlGMFG1qGnCsGNMwotaCFXENqEgctKeDlzz6vd7pcYNKhVRjsaRYKsoCH2KTfsAf/VRutp9k2FZqFhcbpoBGe02ctJpqcnZsPI/502pASLfrWKCR8F47RW8ONkobcwQgHDABTZzRzOitrkXW72S9ZLgVQKJY02fvyj8yNzHmtiQy9IXOz47lt9GnI4KNjdLpUH/K0JlB9UiSBiSxCxdSqOv6XVklWMVajK9VhleZ3WFBM8d82hDp9DPYAjcyMJmtUgVhXIRRNYAH+8YRTqrihS3TD9sapr7cj84TXHRGzdv3oelqTNa33NNR9Udeafoupyieoqoaz5rG+P65uqtaHRqVXEBy4PHV8+hvEubyTbOyncEUbBFLvnoe30vtV93wp/cJU/uNLXNrBNiZbQ/qx7YZYnnopWZ956s+o8MMx7V6f556tOc53czQdX573JaqpTrLrcbEkruCbXbebj39/20GZYICGl4Y+sAhptSjkVVEqqursH64viZW0yJI3KwgFnIeHR4Uoh/PZBFqXcqtLkE9rXnopmJwcBfpbABgytxPtkfX5861MdCEviJjoZZWGbBo3bNITAzf2rYGvlpL73OCzVS3Rij50oib+5eYKgmWNO5SI6XCi//mnxC1Sb3TDa+hQQbHcjRI9DdPBvGxRoI2SXQ1KMzkMw2vi3+2GY8LeH+pJ7n2KM2Q+u0I8HgU4cqCWxsTIr+EvDMj5T4pG59ZRhJ9GSbNDCgEydxL8i2RTiQT7DmcLMnSmY20tPmoJQbTYhXoOzvhlOC/tj2hXXWvFi5ifVSkUx9RhSOll5UFAsnqIzGhD0bIGwkKeKdyIRPhX7eJTyIqIrJM4isiBOEp+FE5ytXs5m0jiuCfxJZEdGve6rYx5v4XEC3DSM2PcI54FzyVgWAoIFGQLgj3SZDpcOwNZk4ExytOsuFkblqditaIJrTGQKW8Z2oioZr9fly4ON98mVgirZau0D2hcWrIN8BOGVk48ihdoSn0stBMUYSttTbb3VUhxFC6ZS2eAn9OsnSvFO9LZf7uRpEzQG4WlJcTUVLJZf2/GIZpEY7GLAnyyNj2lpaNeKbC2G7cVbMpWCZgSbinsUeM8Dz3dJXFUj2T77jg8x6AlFVRYgpR36Ce8SaCecziUBkgE9gQ471SxIOamgHVM5pZYMjDokcmYVUdloHq3FMj5PESijFW3sqUhStdhT1rNFM7yYJxzVWflx5fRUHvhS4EODQ8WX/uKM9kMVZapHYPfdF7ca8Lt9+TMJPPj4Tu+N84wP7psvw9NR3YRjNHV54coknK9az/c29cEqT500K0yIFNmjiPcoUaU6d+JL4hnhMqlVHsTm2SwML5eLArfp/I3x1VfGZxzXnwZhhA4i7S3X1E45X8XiMFWNl+M4ARuNtt1MBRlvZ3SuunKquIsppzJqmedY4HAZXAbhdSCxV2I1CHhiZmXj0jayxyaoEz1ad/pm20IUfZs92tx0MpoPaQFFvbpXtbEsFda3rbWNdba+hZTU112ZIyJmd/CXo4riVIQz0iJRFEbVys+MF7Hvr/BxRLuyi68CYBlchH6Q+h7SAtNHUMmvF8QtzNKegfHn3PwcJP7sF7bgu+oRjP3R5ckNw6FgqeyE8YsuD/dh1Ll5OY5JdEUX7/B1sCSKCcVLk6pVsHEjn8TZQua0vvj3E/O05QuIrzB7UECzOPrwpdgvnOgSl90fGgK/rd+XJLp5Dfasm4TRo9msWlGXbouiZ4WritMSqStEZthglMxkFQKQVkTm4RWp1nSqXJRQy/NjKESAtodamQ3j9i5bKwzFiJNHATgOyOCzyJkTYRF4CfGQPYj1l5oy+pXd46U/8x7xdebP/OkyUirfpdGStNSbVEVm8Gxb6jKLomrek7/UeJqELIiNgR58E1e70lX13zgx6XVyIOFjAfYJCxVJoPSbCEnHrRcwNKmUlQQRB3cgYNDAo8vQObzwUYRdpc1XAM2/iZA3GsgbLWQ8gzHA04ArCSKOHEXLUVxlJYhg5for6AmPSYTLQHIE6bMiySV5Rm1snH94LhjAklTLgEY7eT+FzZjWV77cSqxSGC0nbJuBMFJSANanqJseBC0prqJQCPvBRrJsV0Ip0RQzOUffBPvVp2x8WAbxcrEIowTMII+O/LXicnWqd2doLMmZsn0iilbWRKFlkgtn6OTypfD4wV1GEfTW8seYLKgTY4pejBJKytbq5irbSoMPN+qnmrLGDDNg4ybPP4/9Z/wgq/X0PVvZT/mXPt6p5ZRLnISJM3sczmKl3OHsV9woRctpneYJrDgg1YIbpwigUG7tOroYAAAnBxZCZzyOiStyIaUF1hPYVPSRmlgpWM4RE1lumohly5+/BpJ1KmLIC1f6B9I2iqzkjFqGJylcJrlX4bUiOdDcb4k/PU8KorvZVnQ39xHdzQeI7ma96Hjh8ucNosuLLsgOEWslqsiHHBfjaq+xB35JrVvM9PauXD5KZ72FkBQMQVLFzE+onAKwPk6FPUA7KteuA55bLE7dVb3QXc6h4bXciDgJeToj+FatMNBKto2KvrboRkRc+52rZt2wcXtEuvpQAj+nks3gaX3o4QWhPqa4OE9KVonEK6fKZ1EhtVqxPcqippPgK67p3pTv+cJsKU6cx+zB/MThDatQNo2EWZB05HyVL4AXQ87rpq606/O0KJxn2Xrl6/whwwL8XNgOICzuL6wO56v71ZXhhflltiaG7gB5NI51OWaJ8sKEDNFZrUGkiQWG9RV1yORbnMsoVANf0E6XHKNAcMk9Lt5WhaOfv+cGJF0A/yIluWZ7U/6j47uBuTZyETZyoTQMs2V1a6NtiyPxhB/ruIFvL9+XtzEYzufLcJOQRnX8QFOF99EztvloLm7U2VQos9XdokbW1HATi0NrGDllbxuCMDua3R1gvD9Xd1AKYV3aAYNPqfTA0n6LjET9UPQQxI746/Ud9UjlkLsQ+XgimCOogaxvMpU1FTc5/M16eN6N5UO/gduqZwR65wg3Ld0K8yfcKmCw2ddScOlRcBAAnxrOJ5m8TotKlwYaDyVrmHcAxl/o/I2xb3z2WZ5cQk+zkl2MIChjyz1WuO8U5u0kHZV0UP/IjCnvyglc8jhcBomoew81qrhFlOoW2Dd1oXIzAwe/5nYOBUsNIhm2VKslg01W37xHEPk4OhQMNZwuLszDKq+ybIrFkIupyE7mvbmRd2iaEncqLZQZUKmWCK22Hr9k4ll5HYO9dVk+dXyndLeK7hxm4tCuR1ByK+QljhvGTnHVhtQPrMrV9iFOFPV/1qgsdlH1w9ynKdXXUm1drddWFOdK1tXVOl2VXlYb9HS1TktLdXSl1dFVuY6hkFBD9VKqrUPWqufW6iKZrytVKVfrlHKnNIeywVr+o+nCWzM/IL9yp8QarQGMkyi8JCWT+WWUHzsLBI5/XzoR2Qj9XUhNrcocJ3grf+gQu7N2YOOF1a05YRPp2VaxJppkDU3vtwkmS/1a1EHr1PjLX9BMxbUJazCEfrUMpWxI1Y2jVuk4an0aRz+No2vG0aOPN47ubDd4WiWDp/Vp8PxTD55HH2HwpL/FYFh+ogpJpDhFNSDX6bO8XklIwAFdF+So8anslIWySFtO6Z5RMlxysjkyJkbuCpO0akQwkwRbmnCM01p43sFN2TYH6icLJxLR4AMNwKQDGZ7yYY00O+xFrJsSLFkguM4Zj69BfP6IX4/EhpzH+PNzlE4Bckfubo+kIUqIbgtYmtNwxFjmeiss3TgXL2eJEPP+aJGVdHKYEcnj/9BAtwqnGJw3OSRO6Y3uGb6RkO7kKApqO6bz86YKZyn4cvwmlz7280fMWWo2iwVfN7mUwVBmjc3dPA/qhd5ytoxleHr6U4ZD30Q8tXz565cpvVFhu2sYecoxXiLiHue7sEw/P4cMCdQ1cz5pQ8F03j4Y0lE+WSR3lxq1FvvfP15FoOWJIoRXrYrcPFxFbrZSkY32qqojIsJaJSmW8GFKIiH+eZQkPWdNjY82DG2Ys8HZyYOWp2snKPMjMTGeL5wTlC/BwtBl4D0Gh9MrnQb0/KtKTT5L0Q8oVWFeTp1RYRnfnz6lLKGxGb90gi8nTQ/wzGYjK8LBqJWRFjSbidwCFo8GfkZPBqbefHo2cAl0EjlBjIvNX0b+lAUAknABfcCkjD4MXMdRuCBRclOt+HNnSpoRQe33gyke8ULX3ICQvEptCwLNZnoAafN9GM6RgLUlIp421qQHIcdgmVDMzmJVScXNqkMW9TvXmbnVKyeqKvmi1fzFrThNfLdYpbsCRUpZTWxHioGX0EqPp6UbCVBmYB5UUmVh+FtXkgguyYmek4zC0bFaKU4+q0LnZ/z+AHllpf7ittDx39EDHFtdMsd+FsqbFriE3ptwIZBbbUWusIphNvvBGZOZuLiDzyrR78fK1gc+tuD8+xzaBYWhITU6dT7DV2n+nH5JNQgPv3qNDQmlvxDiaSIUOyn5G3raNsKZLQtKIJyZrMVCHWZHHwldUYWKBo9ts+i0bNVs9bBzLVZiDTWOJchNopZuR7n749cM0fIou6rX95o6SYgtIj2yWSuymaiO5Y2Qey1C+1Pri9cUI2SZ5pegb1/c+qh/VP30eLw/EUq7iRPVdM0jlBrmMAr6mHXyqRNBE3SsIOwjPIEdZSb0CIZW0bmWZx5tYciS07MWxlYXiYl3f/haqn+xSkm9ZhQmQB67zebQ9Mi0oqWt6Ze5RkXY7WuziaTxY2uN02lq03i45t9bg4uG9VoV1iinLJMttHNHtqBBVX94sErk6BtblwC6hQoVoMd5718OlGrBOhjo+9O5oQrWdzmkpLqZUjZDqpWoxJJWKqostv6c8kO7CA0FYTBmMLcUaF8oRwN3WuEnSvuuVrYCkAZevglXJC634D/IR8gzaLkzJ45/8OOkBRYLWLrBJERR8msYMsvpPuEhDsZOOVy7TJR7QWxV8WvOlKL4ykyCwPo2xcZrSSrCKc5KXpxG9d1/BF/c5m3k7uSdsoAHL8AosFaWKb05Q2iVDJlHUY0Ku0yjoqSmS3GKUuIAbI1YebKzksVtKQBxQhbi3IZeFEycDGWT1CqnRkWRE9ObB8qJIWdyoveiVJTEcjFxgDIxpcmlYuIAqpiyz0UPC+3Vjume4+kqMR8g8hFinXwZ0aLsXAycQspj6G/wohJ6XQW5SiSDvSALAGhBUackYZ/y0IeiW+WAO/JMdfmyZX7EhEpJat5CXKgk/q9WOfQ9T/HgBuyICHSXqW40dEJR29VDkIudLkZxZkDhI3a5nGLe/6RfWnQhXSWzP+ZsQ1ylsHcdj/qGrhY38uOhNq0u3VNv4SE04gyBXIKXCz6MS/xneW/DekhJiB2nTDzTqOzccHG/DJpTxY5gy24gR9Z2BEIya+kVS5NCG3nF1iTx5l1xgptiYloqzYYqkQRJHiVJ5I9hMK9WaA03xKpVAnGTEG/a+ngBPk5RO3IrIOud+hQq7c3/FuLlWU2VCE/OBaBPr9SpP093Jk/CWjGOkV3TVNrAikbYVgJJCTOJ/AgdFZYz/VoZqTwoA/nH4wE7oyIPeB6WAnS/cF3R/VAIbR2tK1LiFlh2idNHjEOnNMtMzAxADs8Jx6yvXtB73WDgUHo6DADhwika68f6gnHoG7zSyw+mj2c+MPNK2hQsTPkE4Ba9SsX1xW1KibsqzYw0jb2AoN5pZ4TSGIBg85esBMlM3DRMJ+DIUzwcpEV9BMWByrjWwIuxVr2LhNIWN2pjKaidy2WokNogTDaFk+Gn2neUi459UQ7f2sSuYufL3LNARZUGKWoGTcLge7Vptrpof1ktm8yVQyI+pISGrCc8fltQllwIXFtw45R7Lp6zp1mL8mDSphw4UVeVSHo+Wc7Swqe012ly7t9u0mNONEPA5Twb5SlwI++amiT7+c0dKRQtciPNjL5JF3Ik4UKHBZ9zJHiRcKhc9/PdsxkWTcjx6KuEmd4qWURlKTkuexeQ7zT3s6ybdhD0gIZ+m4YkErE7WjfhUAh3sQsxaQQuywCnVgT68CrFt0roC12nLEGx64z4bNAmYtkMgkiNB3EEcuxLTu9O2ZWOAxNQQ/OFLhmK83UtTGuXCxhmCKY/i8L5a26Z6rcRSqsD8LQT73jFjs/g8+H0a/XaD7zwuuWRK/AwaK4USDqcvwQGb7uSFveo2VjGnpgVvm7KTr7/Z5EBPPIulnEyZzFAJaNSquU37rhxfLx6Q90+eklJFLP9/FXVhpbkIGIpZZcJWmpxEFtC/1opRM34UrhS47BwwMmHTenqj8C8F01nmYSVkmqic5aoh2IJxUVoX2trsqxRbZhwhjZGMxQOrhItLGkKjfF37STuuaC3aovhi1KI58sHDXO1miP6C0ytvqtGhAaU6S2lX9yWKN6dt8AuSAmLzLhbX8jfYHnzU2OUEIB7Tp3GRoYujEa6/kAcrZSCC3GL9IFlXAw8lOV6V3Aiy4MWciHXdG7ZAXIzj80zC+TZaU7wOQVdX+YdqeZiNwpns28km+yWDrq6HPA6Q8pBwxiTc+fKx4tcK3jEkhMklTs1D9GRpvk8D5LwF59cV2816EBzFrqX8CUgDihRTvBO3vRelOc8XMZUM1CmagwtPfwI15mC5MQFpxjaYoX7rWHkL2+lPi3F1OyAztxR1iz5Zb008BsGRDhBUvHiygDLDoPOdkOeXSnHJGaLj7KSbZMpFXQ5e2wuWrLSN7ibdNEs3RBdModWyCNRrJfNWdANH5uzwOqQJ8T0J09KeZcLeM2RObXRFoqgiloIToBpBBA9NQ0tJuKlE2UZZaaa7HM9I6BH/jaNEijY/LsOnR4pljcHaOkiH0eGaXz1lSQyEbiuADPXS+BYdhAVafGlDxXpvFsFhnvHiue43dSxXKy06IruaJw3LZtaDrbKul6WdVlhJYncldTU26ymvk2d8zVV9TarqgxarKtvNd68whttt+vrKp0v/sDKevshlVXsXu5RV2/vUVf55HildAvIVsNXyCwCdfTacojZaoDZyAiN020aQ1+lWw0eOJKm+MXhtLgrBgEePjie5aMiy3Sknqh6dlU8AXY9F/qTPbIAqDzyVB5jAvH2MViuaOOG85DVwhSOO9H9bDdYGfKNWfotRfKOqk0yKdmIVUBT7wXQnxCyla5enxMy0+lq2j86swTyLQa2PDJLnLfoYjDcWot9yblCbH68/BMycUBx1JBofhd1wWiv4d0LeTqjrdx7s9bSF2VSBMR17Qq73DWAUQPhfgC+Z1u67Nn5dCKiRGikboQucnRkdDaVrnlodFT9k3I9KGEXT7MpHGcjsyu91uUdjLqtfqXlODCaGwtS31SQo7KC+MG9CtLcXBDNZKZIYnsfuPjwsKmXolN5q0ZyacN0Ke5bMe6HIVxjT9Hhr5U+r6pISIKW4rgs2Cxk9ps63fMxc9P5snl9aY6rXThx7F+RfXaBy500J6ZbRbGuArUBjMKBtSxQ9TIgT/i1PR/vuFr5dCHxoFHN4lcWvOBA6RcNXBYuFkHZx5F8PFPxCL3tD9ErHKMnH4JXPDavkM5YSNTVYuUH4xXPRaIXSrB7G/jgj2dNgJbgkoh94aIl0Bw+6io7Y7ONRQ9YRVq6hcdynUWlDCrfvbMWbE4jjqzOK2ar3yXzUljBYvUDPL2jKXvOhXkF/VYTPTAax5MZ3QxcuWIqXQq7freRRmKl+yhUSCxV5plX7Fa3slV8d1P4uqzMywQzpJW0WClbNnbW7Q3LNnlJjWCtfoggZcohwpQsQN5RJp3f4LBQOMYcU57OCr74lhpfVuqUrGgc8U9q2W1x7XIBTLv5xcCgPDYBcDWz6UyLrvOCkSgtbHrSwB2Zvyulv3bxdgE6Da40u+t4zr1Yy1oHJy1/d0kgHqxTzFrSd0sHKLtL6oBwp+oDGwH+PCrByvtv0woe6/m3q0Q68MtaUbZkKavL4hn+fI/vLb8sQLGJ9oufmB2nOU5/X/cRoe9GO3c1aoUl5z41AV6FYfJj6JFqrXUexgl4m5MgbqUBqPRERngcHeyB5ewvkqODvSQiJHZhCGhGy6B5TiJydIAL3A26wupwdxLOPLzN4ywAyrtHB1RORwd0YslAy+Fw1z0n7iWUYVeLc5aE0+kMUfcokkyenvZx5ozH4AXvHn01S0ZiciUIE5pYObpwVi0qBAOYBwiJDProQQ6YgvDD46sWmDztWglh45//+D9myzT+8//j75N//uP/0Xt+//mP/wt/T433JAr32+aGTNPk9A+XklhSCk1W8MUj3q6chvvEQRG9MxcVC5RjbTok4upPnpZW6lmmjgqycrvG7tGrdIRnCtFqtVK+tcpANYVXNaj/zHdpU9gL3YQkzRhwnPnuUXZlFmtOVCfZG6rkKAuJTOg1AWtVdsTAWsDD6zAM0ms1OIlkwYMO2NtKl4lUKwmZL9BcQTokikBgYBjihpPchGW3Inz3+uWPLepAV5Fgi+8LV+mxsldqUl8B1BAlu1fEkNtZS3cnBmM3v1AD+4oNjTAtS1qNUpF219TLRRyC/tzu4prI3X1j9wkLu2UH64DXgMFC8JMNhy4AqLWMb9Fw3WPfcSkXuw2BBhp3G8aucAcCUnz03/wnY5le94Ac854AE5SrFWh5fnr6zTb/I7pw1gOg3u46FvyxMEMTHtrgO+3me33hy8ntLm5zwGwAFFJp10BfAcqwAAe/0j0O8BV6nl1onpTUKTzflNEyi7RMYMTUE7NMSk3ZXoUk8bN85BItViq4/HAmJAyGA5LF/nF37qyyZx7ZhXfmUu5m52BnINlx2NmX9C5o5OJk2GNbENr0klqr3WGvHYu+9vqQ0sP/6esAHvt43S57HQ4A2MR7Adldvia+t4fwy2JX91r4oY+XT3bZBxs+2GYbPgzZJbkd3AyBOdhtBmHauB2C/mKZ9BB9iHcVmib9YOGmib5Jv+KHfhdZxpwGFKONj3Y7e0diw04Gjjn2sNA9eg1vG6h1kKJlnZ6iEknbBUBEfaa6ZVF3ph1prfGbBndTHKoi2CB2MSq0FuxEBqEnmND2BEqTK15+dxhVrvRLpmybcoHnDdlYhWzaci7tjZmcbi7tfwQGjAYyHEu8O70TO2jeGcMjDL2F4b5gGyh2zTpD4O/GE7pkad94fPyzkRoaX02TkT4XxgAaANTShYGBBTT2DZyA2zV8Dw2w6CwbL6y+63V6ZOx4/X6nZ06GlmWDVhO7O3AHFuly9uA30jtiVwWBNTFHD+ZM4nYfzKZqSzLv9s+duHrE7J2WzvCjOPvURiQeHrhQo0P5DMOYnzv0B4ZDKB3N/Y+xPQ2e/TY2qCZJKwtdrQL4zFnEmEhNWY1iKDQ+15qz/9PM3rWCKNfkP9hKpnKrnpyYLbyCAf/Pr6v9n/gRxgl6ndQhr+BaiVg/rLvwvE7fHk/I2HXMjjv2Bh3Ta9tWd9LtWvZwbOXdxf9CL9He2GK6bdpk2mmTwXEcGo2FjQYkSttN75O7+Mld/OQupu7Vm8Xi5vi985RcDNznv4f1n49/uzq+iC6fvqlPyevLuvvdm+nx1L98/GQYvqn/VJ98+9v0ePF9929vBr8+ux7sdabh5fHw9x//djH46WIe13/84dn5cXvqWxfd6x9+HtQ73/u/Hb8Jf3efDF4899/uvXnxXfe44397cdFdvlz9vDedJwD/6LdXb+rh334a1H8Op+fH9e/ePnqyZ7fP3b1fXwTL4+Sm7rwfvPzt4nrv6Q/R4+NJZ+xc1H+f3Jh7v4TP58ftb1aXbzqPnl4Mhi/++tvvxz9F3V/f199efz+tf//d0/i43XtE3gx/erKI6/7T/sXx/Pn45fu6c/746d7xCyjy5aNX0Xv78rg92Hv/+8A+vn5sfvu+f/1s+qLeefQ2OH7t3Hjv934aJm/rnePxT8dvp9d/e9P79eUvL/b+KHf2g9xeO3V7LW6aU/+3t9b/fZDPWuoB20Vq9hpvGk38DZ64rce1P6Lv3DRbw2Hf7HUG9mDQt3sD9FszX1qTtMm17pj9/gB+Dwbdod21hwVXW0Pzk+f9yfO+j+dty7nYHyET+5N7/5HsdXtgu8OBaQ/G7rAz9Drjge31LLPdNcfmhJDOv99e/9/lUf+LXYT/3q41uKfDTr/X6xsGuqjD/qDbHw7pc7tt20PolnNfFp3ZTm/Q6w06PYTo9oYDGJUotNXrds222QYnN7tPA6l32oNe2+5SmG673YcRjj53Ot2hNWjbCvX2oD3stdsIMRj22z17wDC7g2GnZxgq9R6MnkPLpNwMTbNn2qwcdts02x0YECTqgwGMRaZFoTvWwG5bfZoTDFngIMMQqFA3+x2za3f7lHp7YHU6bUodRuNhFyQmUx/2uz3b4vn3B+0+/FDeBzYiF3gfDMxup8ek16eSZJKxOsB7z27L1HttczjgkumaVhvy6tDnYdvudGGcVOXe67X7QJTyDuP6wDIpbr/dHaAMZOrmsGtC+boUGso5GPZYSaEMNpglBbl3zY7ds2n57EHH7A26NCcTsu0MTIV3CygifVqrYOX0BqxW2/awO+y0C3LvgdwHXAuH/U7P6nGNhE6y0+72FZ3ptIfdQddmpYOS9jjvIDEbyqFSt4dgTrVZrZqDLlhWFsXldpZSq30L6sI0++wS0O6wZ5q0pJY9MMEwGhZ0httzlGKPqQTTt+Gga/WV1mSZYEu125SibYOW9IcUc2B3TSh55/RDYkYf7N2zs7yee/Qg4dwhte0esSeuRTzX6wzdydByxt3OpO/YpuOSSYe6G1sHr0XgjaErEXjjuHk64kGI0I3CMNGGIdJIAcqKA3FwJTaw28rEeYbJuwXUx7OQLqJL31sufqDZ0DvBRnlKRMAicMmvfnJeldAzopPImfId+srKpCf89RmHwDKk0NLqlpQwpmsO98gqVzncg902nu6+zyjDK18T9c3NcyCeYuO1OyN2T324jFzyhB5rVSLDz9Hw2TXqhoLOD40TxZJRa038KE7zpiUDhDw1D8vcyTWhFfKmKMx/AUcQnUA=</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_5aed62c96b39448cb9c45352442f166e\")) .filter((elt) => !elt.dataset['step1']) )[0]; root.dataset['step1'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def create_weights(seed: jax.Array):\n",
    "  return Weights(\n",
    "    kernel=random.uniform(random.key(seed), (2, 3)),\n",
    "    bias=jnp.zeros((3,)),\n",
    "  )\n",
    "\n",
    "seeds = jnp.arange(10)\n",
    "weights = nnx.vmap(create_weights)(seeds)\n",
    "nnx.display(weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fac3dca9",
   "metadata": {},
   "source": [
    "## Transforming methods\n",
    "\n",
    "Methods in Python are just functions that take the instance as the first argument, this means that you can decorate methods from `Module` and other Flax NNX subtypes. For example, we can refactor `Weights` from the previous example and decorate `__init__` with `vmap` to do the work of `create_weights`, and add a `__call__` method and decorate it with `@nnx.vmap` to do the work of `vector_dot`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5d9a55fd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y.shape = (3, 10)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<script> (()=>{ if (customElements.get('treescope-container') === undefined) { class TreescopeContainer extends HTMLElement { constructor() { super(); this.attachShadow({mode: \"open\"}); this.defns = {}; this.state = {}; } } customElements.define(\"treescope-container\", TreescopeContainer); } if (customElements.get('treescope-run-here') === undefined) { class RunHere extends HTMLElement { constructor() { super() } connectedCallback() { const run = child => { const fn = new Function(child.textContent); child.textContent = \"\"; fn.call(this); this.remove(); }; const child = this.querySelector(\"script\"); if (child) { run(child); } else { new MutationObserver(()=>{ run(this.querySelector(\"script\")); }).observe(this, {childList: true}); } } } customElements.define(\"treescope-run-here\", RunHere); } })(); </script> <treescope-container class=\"treescope_out_3afdf582c2b64763a175ed6a40527896\" style=\"display:block\"></treescope-container> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_3afdf582c2b64763a175ed6a40527896\")) .filter((elt) => !elt.dataset.setup) )[0]; root.dataset.setup = 1; const msg = document.createElement(\"span\"); msg.style = \"color: #cccccc; font-family: monospace;\"; msg.textContent = \"(Loading...)\"; root.state.loadingMsg = msg; root.shadowRoot.appendChild(msg); root.state.chain = new Promise((resolve, reject) => { const observer = new IntersectionObserver((entries) => { for (const entry of entries) { if (entry.isIntersecting) { resolve(); observer.disconnect(); return; } } }, {rootMargin: \"1000px\"}); window.setTimeout(() => { observer.observe(root); }, 0); }); root.state.deferring = false; const _insertNode = (node) => { for (let oldScript of node.querySelectorAll(\"script\")) { let newScript = document.createElement(\"script\"); newScript.type = oldScript.type; newScript.textContent = oldScript.textContent; oldScript.parentNode.replaceChild(newScript, oldScript); } if (root.state.loadingMsg) { root.state.loadingMsg.remove(); root.state.loadingMsg = null; } root.shadowRoot.appendChild(node); }; root.defns.insertContent = ((contentNode, compressed) => { if (compressed) { root.state.deferring = true; } if (root.state.deferring) { root.state.chain = (async () => { await root.state.chain; if (compressed) { const encoded = contentNode.textContent; const blob = new Blob([ Uint8Array.from(atob(encoded), (m) => m.codePointAt(0)) ]); const reader = blob.stream().pipeThrough( new DecompressionStream(\"deflate\") ).pipeThrough( new TextDecoderStream(\"utf-8\") ).getReader(); const parts = []; while (true) { const step = await reader.read(); if (step.done) { break; } parts.push(step.value); } const tpl = document.createElement('template'); tpl.innerHTML = parts.join(\"\"); _insertNode(tpl.content); } else { _insertNode(contentNode.content); } })(); } else { _insertNode(contentNode.content); } }); </script></treescope-run-here><div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtWgtX2zgW/itqenZIFmKcNwmFsw7kRQsUQgtlZ05GtmVbxLGNrSSEOfz3vZKdhxMnhS1td3YGziEgXem+r74r8S5gE5scSswnJNBcj/R812XoD+S5AWXUdWrIJzZmdET2keE6LGvgAbUnNTRwHTfwsAbjY4sykhV/1JDnw4hNA5YVW2fZxINRx3VgWMVa3/TdoaNnNdd2/Vq4dB9Ff6k2EMB+VGdWDRmUAZnDiMP20YA62Wg8J8v/gL3ch2xAH6ljwjrX14mfhaF95GFdh8GsTQxWQ3nN4tI4JGsRalowkpNKnJ/DMAXlZvtHv2RHNKAqtSkDFfGQuTPaLHWYT52AapwtCWcjvZ7e7YZ2fDezY9YfOsDTh7FA86nHEDfEwRb2PJtqmJt219UY4WbyCR5sHabTmYNDsDzwCxjSieEE6AAxiwaSSdgluOXM1Uk6I1luwCQxD6oRhnoecbjKisZ35Yv+/VvSTBs7uk1g2hna9n7IQQIxu67rwGh67Pr9DFqUwb2GIT4VG2ZU44Me8Q3XH2BHI5LjjtMZEQjAIL0yg7LhoneokM/APtRA6SWpJZs4JrPQwQGSOclG0X3Chr4DdkfEDshcMGvocMmWtw4sajAunyDgvzzB9xoOaQg/R3fHkk/uhyRgikMHwl1NHw9IOrRJhu+xv8LIGwZWaMb9BB2nLA5CNTZo+XwZuBShI5lrmnaYvj2RYhCtHt+LjxCb7SAyggCPPMmlE39LfTLhRk/5KS5QRCxpNg6CD5DF0b7p1GzP3gDCMDVl/pQBe0L4ixg/fLeblAA6HSGx4UEqXmdSiGEVNCUPByk5Banrs1US1wERwRgOTG1KhmQLpPmaqe4pSMaw3qkuY+6AF4aa47K0ZLi2jlVY7cC2NQsH6UMbq8Q+jM/0Qh5iTU2ziNYneiaD/slNNy08zPVqSJZyJTJYLT187DErVIa/9xPL7JMkymEPq6pPRiK6RXV8W97LY1meE2juYABqLVBg8cVds0SCQ5Etd0T8TAJ9SE70Hhd4vmFUkFfr9gj76WxWtV2tHw5l9qdVWOif8x5Q4NpU30QZOuFrxE/cvcTnwnkgDrHAIcRfVLrCv6PjSfi3hijDUGb54pj/1hxsYK4kNwO1TgNgOpkeYMuE6BCJMKnVVAIlb8FybzXxlezh8GzK5vjhFB1q4NYZL+qIsBFm2MBTOHOVs479fkCwCWnkrK5+pXCfycCXJi+a0sckFIdsDW39mi+p2tbPFC++aK2Q5R8gJPcjZzz0A+5AzwWoQfwEvjR4PbYiFQSjrCiOwboYfx2uc/UYeWCrXCQa9AzqB6znOmERWk2tTakk5Us8mxJdhb5Z/NDjyyJyrQbYNwEchmKIhH76Rm5Qir2JOoTS6CQWoPl0UtCmUGqJCgwJ0DyZ+FeSK+qpJWS/dYohKii2UXcyUF07QOdDxvXV0VG4Ej69CSRGdkzUPqD0sPIO4OC1BB7HDoPlFAdEn2H7t0Tm3/urYR6uFphalqr8hIxrGeZHghbJ5W6+UhrjoKfBOQCGna3HBoudHtM6vYnn0po4y0XTRweYjhnOYgccK1BbZnGYM+FQ1MfONJrFtigXIAIWA3CQdYfsZarMJADHUKK/iUsiWKI3dOC5PsPOyt6q7/aJ0+Mj82L0desuLFuw59TNT5KlWT3RFPZ4diwgpDBZZCnM2AU6n+OkBUI/xE0zSo6bQFW9p0GjofvEiZSPt3IgZZzw1SDeLBmjozmW/Bq2tTT0m9DsAKAR+FgKGObrZ/J+N0kiRBVKogOugjDhsGrBHfdDbDvQPvSgQzfoA2wSS7w9kXjQK2GOtMbYd8ADvelRMfWuYWAtV0gg9KDZWHGcqIeRkaKhbOTMOaSsCZCM/azpY52C29IoVyjpxNxBLiSJSZAM4pU1aydMGugSeAkSQygy84osK7X6dQ4CtFLys7PgXPb1uthcpptB79CHkN1xkh8evbJU1KzMc2T9MeEctgZCqsQGIdZ4CT3AJ7YF3krnQJ3FJdG+AVCRDG+gIxvMl2afv3bKd9rGiJWzLaej2Wj4B1pz3ogJC2Z5Xg2DSMEcGYB26xVLwFAvRjBfC/E/Yr1i6F9erZ7RK24mfpJgx+/S1c+d/C2N/diCA5P39UnIjgNtMF4iTdgSLlCBrOs2CncAk9jYAwT29S725Q5ez2EuaEhEHgCq6GtpXkOOJBYxU8yvw5JMEb9vkuL3XGjDDsuqOrGbt2TC5zBbu88rXjbz+z/0RvF9PJEM3x2kdVcb8psiiYOCQBphe0gg0DJS4A5IWkAFfnPIP6WwUeC3hs9sFVJbcHBmZve0gUUI45e5ZIyOut0u16bLx/jVrJiUfCLueroTR0v//q+oPdHIFLS8vFVZvBxy+OW0HY2Noxwu8pu1wNdqaOjbaY6ba3x+d+waRn5fBWReLu7ocrV1aip1RXx1LhTFFb/VL8fws91UlIay6as+UBSz777XO4360fiLolx9OTpRTjv1I6VpPnTaHywW1E8pMQvN45v8h075y6jrDenH09JV7uSmc/n5dHR9+sg+TprNo+1rs39F68eyRY8vhicNvXUnt9VdY9TRvfv3Zev+mtKL4anTstrGJ6Z8KtfP/KLS7Dj9Rln7NBw625eley3oj0dG0969fzAb7p6pnoxbe7m2susol6UPvn+Su9w2H+VLXVZOjJx5Vjkat+7ypuxOhpeVyqCRK4/bN9Vz0/TIVX9SJB31saSp/nmLYcW86JyNj3EwCS6Gnc7NdaM5Vj5eeJ0v+qfd3W2zclW5KTDZeP/xXhmVYM8PyllFOR0rA/Pxsrs9vO2Sxs1D3ihrj2fFy/akNKwr7x/rd17TK9D2xVFDvh1+LHYrjlH/0Gg3TwcK3d4bNfKWk7Mq2+rn8c3duO2Pjlufjpw7o9Ew2fa5dmvblVL16GRc37OqxdPTVrfQulXMQad0V7+osqsWaVcb9XqnVTg2i5e7X7SJqrTAp5/f7yoXLayQ0yNbaT82zs1bZpbrH83z885xvU8vSqRZvzmqNzUqe5bveg7EhnfbOM495vpd48hg1uS909ZxM2gb8tmg1Tgr13Xl/vNnD7OgezvQdUyreeOxWvxE7+7L3sAvn7tfjrrUbw1GJ61C97pbaDbyWv3CuNpu267XKjaDcQmb9+U9eku6Z7Z37dTbHaKf+mR4fd86GuSum36/230o5cvX18FYAYkySDzMsPSWCOstfmT+Dj9m2Y9114PeYZ6S4jlJkqQNFDthzv4Ge22+oLfE+4ZoGMNeFvaG8HA0lA5byvjrE6TglcvTF8iilpOPBVAe+Ba8h+aNJx5jypCDR9TEzPUl2NlTXezr0tinjFyRB5ae78URRbjX/IkDjvh0aqHB5o8bwOWKDgh04unp69fKOp8MoFteWfq0g/KyLAskBcUXUGtaXBUl813oolNz4fgl2bSC8fegFHqLmpjaUNiYizjxG1HZAG060NdBNaZgM4J1fgmwvWi76KHmK080/EIBifJ4kIqBqhpy+7ZmAaSu7JU4+pELKF+SquVybq9cyheKhUoJ7QICAn2TACW/2U5F+0dvQPHb7OVGDYhDBPCOOt4wOslS4sxX3YdU4iYRPIDJEBqAkmJxnG/siSPCiSh+SbEkafzsTx3+YjNuSKDYTBebXOq7YTbaIfy4FudPlyPm9GwmssAKjynUSSVoNn3rSR1CrHzEPga0WpVRulCWUT0TZ7q4emsKJbfiw3yoZ4/s+HBqpc2NY/vUBuoVUpXi4GB9bHxLKCSHwPdxqGHjB8lxHgA1ARYBSWyqSnGDC3+kkxy7Nb9R20Kuc8TryMHWCwupeOXMbKHZdd9BSuLGTc1zejYFqEnM8fvb5ZtKmBNFAw4lC36PFD9cHzwvDMsCxGMu/6rxuHInmNo0nzoUyPbgGd6m+kFq9vJYxHsFA1eral6vFuWyWs1V1Som5cKeWtGJWlhimvRiGUVbsqZ3EEECjaMki8xaQkE4JTFsF7NCPp2Td1Ah84s5D+WvuCvWYCSqW83vVUtGTjYKeqlYKFVwmZTUPZyraIVcuWjImyr662RtUgn8n7Lq13Nh56X5sqnqfNcyEbZ7m4pFSPE6JWO5LEBd+EwDKKr0UeyILKpDzgKiQfM2mPfH/1Vc/1+Y7/kfawJ8kwPWHYKL/4WSOoQe4BnOXM2XzOYDeIV+Z5V+dtEEwfLCnPoTnq7rg3wzmOoT3yH233DqOzk8NO86l4ezPwNSAbpP54t/QkiFSUEmJJ/P5/YqxaqO90oqfObLahFrGKvyz4ZU+VdHVYZc1EqgaUmuyMWySrAs7+nlak7VqpUCrpb+Iqgq/9cAVmFJ2IwNFmn+Blc/0YR/IYD1mpn1Jz1t/4xuXn+VOf9fdrT0r1SplwfH82HJlFKno8P/ALzZUpI=</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_3afdf582c2b64763a175ed6a40527896\")) .filter((elt) => !elt.dataset['step0']) )[0]; root.dataset['step0'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtPet627aS//0UrHuRVF1MUnfL9vmc22napnGT9JJ6/TmUCMm0JVIlKVuOj/7veY/dB9hX2Ec5T7IzAEgCICjJTnrO7jZuapPEzGAwGAAzg9uB610bUXw7JYe7rhfNp87tvuEHPtk1PPdwdxyE5y4ZkzAk7nnf7vXbY8scN912q9nuOh3SHvYcqztqWp3W2Nw9Oojmjg+/kd5RwwlD5/bae38+CvzY8XwSGnfGzYUXkzrAjQhmFM6c6cBYGTrghuePA0AZw5f62Jl5U+BtFvgBxR4Yo2AahPvG5w79GRgzJ5x4fn0YxHEw2zfMht0ms4Gc4zwk67Pz/PkiPo1v5yCQ0PEnZPcMWLgmYeyNnGndmXoTH7jwXHcKlMbeNCbAwwSoRZBOylbFCCArL74tm4125d6Z7V8E11RQedL3o+cvZkMSAkE/iMv742C0iCpAdhiELgnroeN6i2jfaM6XH0aSPVOmkXxSJ136M+DZ7RvWfGlEwdRzs6Q1uTYigCRhpOrLutqjLMTeHHAkRR4Y8yDyYi+AanOGwMMihm9DZ3Q1CYOF79Y5yzQjHcPDKcACFcd1PX/C9Gp0gWQ9H2qoTq6JH0dJZjeeG1/sQ+3FdWQOkgYGcjaeBjf7xrUXeUNUnHyx3tc93yVLyNk0zfWlHAbLLUsZLOvRheNi1ib9D4tFC1TjH2z4wIuuL1DKV38NW6OpN7pyndi5T41NAwclej4jUeRMiKA9SYteHeyxvuQgDgmJRsGc1MOFX78gIXyLRqE3jw2qmyVnPgceHJTAXjCKSVyPAMeZlY528AeyjWIj4cI4NMrlinF4ZNztGAb8Gy/8EaIaLolI6EErf09+Amn0ythoAMAwQhIvQt+gX4+RTmMcBrOyEwdDAKoZ5RklOGuMApecoCiP47JZqQwAe7VTnM0zEEPctLOMGKvD25hEwOeD8kuIjJE2UvHJjcEzorTKlHxjuBhD385ReAEZziaun/v/FJ5pNvfimLE8JbHxGJVp5sxf/fXRE9DMgVqaCYkfgzJ6/iJYRBS4fO1MF6TG1BAwES0pIVIcOhE5p62hZgTjcURixoc3NhiqcXBomAmGIcBDccwB/8owsy8rg0wjIhA5OjSsAiIiZ40p8SfxhVE37BxpqyETT4gxEY/iKKXIsvzaKOtJW5WBjo8XTnzRALmDzFJilRwXWT5fGhbnR6jpUCnQaZbF2al5hkxZwAIjVzGqnLxRhGRUDYsjirXDMpusy8x6aGaWPrPhuszsh2Zmq5lx/T8Na8akZgzP9I321gfTaXQcjiLPv3hFgHqZ53dFbmmf/zNX+6k3f+GAbofOzQvPZ3/xnZP4qzNP1DKlHsUOjGev0T5xeRZlwIidRabBqNmfedEzz4eBoUyT/vY3pkIwVJWXFWMPEYwDwyL1VoaXFnCZaJaizSkApRWBZYbEvqbEvk5h8IcCTINJOZ9rlWP/HsZQKfxtHtyUlwygZtiVSqrbK0GL07YvydE4lPoATGfyVBIU8aPGMPkXFlaFZ2R1nKmQaamlhKQXpbWcQM2cZTmpd85QZVBQ0oMUIuXyX/SHNTuoLWwzxuGOUvNkOS+nKiDLAJphpt6gOylYUrS06tM8QFRJTnusIcti2VMExwYPECz0tYCrazJp08PWMRBUJ8FSVCclBi3GMmUVSLtXbOpzgr0sajuMu+sa6x6nmdDiLSshoh1NPri2MqVPJJm96lRVJYRydYkfzDwfbIww1WHPLwsqoCu20vVxEVAWhN6utoFK0lAQXao2iSml7mSGtRXIhb9VtQnksiE2NwrEwWs0pF/HIVjcbKyXbbd0EEmGLmmIeRdOhuUv7sKV8cXdBH8NV5V32uEG7fvQicCsmjwsRwEC3REfvJpbgGjYlt2B9hlCF93oWm0bnif4bHZtfB5mnVSGdmRYdi+TPS9MiTpAJa1KJyDUiSlJAlUKCv759MSJwWHywQaC+oB/tzVwE+mndKCESi6jmnrU4IM/BwkIN7LgW7VaUSy0MLgBeA546p0lCpKSu2TkLoEcwKakLkVSfPAJbk4vz8SvkEm8bCD/r8goLqN1cQm8wx+vZlg1webLNHKVUy3GqOtNvJgazyehN3NCrKpTClv6fEx/SjV4tMbd7rBFH8fj7tgk9NEeOaY9oo9ux+7aPfrYb3W6Q5c+9kbtTmtYqnGCpNntjmyaMhwNXZs9Wt0hGY1LAEPFpPL1msAXV+asO8b/KLZDRl3S45wNh13OQ88d9xz+td/rd+jjqD003TZ7bPVH/VbK2bg77LiMHXfoDnuM/T5xHdJOOdtJuRuR6fQ1eFHAUnfAEhSnBTyTsTfJ+Swu9DgvffIY8JMejuoeVDJ1W2oGd2G8CHo2z818GUawlnboiUJw44xCC1rIFYQ2YeCyFFxNRxfldvtLDBpUSoMdjSJBVtAQu5Qb9oD/KoP1NLumQjPXsDhdNIJT2uzlNNXTU7NmZP/OalKCRb9a+YSPgnFWyflxstAbGKGAYQCKPEodzpLa5Efc7JeslxypBEoabbzoB+cH5jxWxIaek7jY8d23+jTkcFCw220qD/hbESg/qBJBxJYgYutMHH8LqyWtGCtXleuxivI6qygmeOaaQx0+h3oAR+ZWEjSrQawqkIsmsAD+eM3I1V1epLph+mNV19qR+cNrjonYun/1PCxJm9f6mqs/qOrMP0XVZRLVVUJR81nfHtc3VWtDo1OriA9cLjq+fAzjXd5AtncSuCOMgil2z0Pb6X2r+74V/uAqf3Clr21gmxItof1Z98IsTjwTrc6s9abVeWCY965O889XneY6uZsPrs57k9VUp1h1mdmSVHBFrtvUx7+/7aHNMEdCSsMfWQU02pRwKqiUVHWrB+uL4mVtMiSN0twBZyHm0eFSLvz2QRal3KqS5FPa156JZicHAX4WwAYMrcT9ZH1+fOtTHQgL4iY6GaVhmxqN29SEwM39q2Br5aS+9zAo1Et0Yk+cMI4e3T5B0NQxp3IRHS6UX/cs/wWqza4ZTX0KCLa9EaLDIVr4twkKtBGyzSEpRushGE382/4wTPjbQX3JvE8xxuz51+jHg0DHDtSS2FiZFfylYRmfKfHIzHpKseNwQTZooU8mTuxdk3QK8SCb4UxgZs4EzO2FK01BqDabEK/BWd8Up4H9Me2KK41oPvXicqmkmHoMKZmsPMgpFk/RGQ0Iej5HWMhTxTuVCJ+JfTxKeR7SFRLnIZkTJ47OgzHOVi+mU2kc1wT+JLIDo1r11DGPt/AoBm5qRuS5hPPAuWQsCwHBnAwB8Ae6TIdLB2ArMnAqOdp15wuj8pTvVjTBNSYyhS1jO1EVjNfr8uXBxvvkSkGVbLX2Ae0Lc9ZBNoLwyslGkVxtic+FFoJiDCXtqbLeasmPojlTqWjwE/r1U6V4p3rbL3PytAkag/CsoLiaChbLr+14RLNIDHYx4E+Wxse0NLRrRbYWw/biLZhKQTOCTcUd++5z3/VGJCqrkWyPfceHCPSEoioLkJIO/ZR3CbQTTuaSAMmAnkCHnWgWpJyW0I4pnVFLBkYdEjrTkqhsNI/GfBFdJAiU0ZI29pQnqVrsCevpohlezFOO6iy9qHR2Jg98CfChwaGiK29+TvuhkjLVI7D77os7DfhqX/5MfBc+vtN74zzjg/vmy/B0VDfhGHVdXrgyCeer1vO9TX2wylMnzXITInn2KOI9SlQqz5zoirhGsIgrpQexeT4NgqvFPMdtMn9jfPWV8RnH9SZ+EKKDSHvLNbVTzFe+OExVo8UwisFGo203VUHG2zmdqy6dKe5iwqmMWuQ55jhc+Fd+cONL7BVYDQKemFnRuLSN7LEJ6kSP1p2+2TYQRd9mjzY3nZTmQ1pAXq/uVW0sS4X1bWttY52tbyEF9bUqckTE7A7+clRSnIpgShokDIOwXPqJ8SL2/SU+jmhXdvFVACyDy8DzE99DWmB6DJX8ek5GuVnaczD+nNuf/Nib/swWfJddgrE/ujy5ZjgULJGdMH7R5eEejDq3L4cRCa/p4h2+DpaEEaF4SVK5DDZu6JEoXcic1Bf/fmqeNTwB8RVmDwpo5kcfvhT7hRNe4bL7Q0Pgt/H7goS3r8GeHcVBeDydlkvq0m1R9KxwZXFaInGFyBQbjJKZrEIA0gjJLLgm5YpOlfMSarheBIXw0fZQK7Nm3K3StcJQjCg+9sFxQAafhc6MCIvAC4gH7EGsv8SU0a/sHi68qXvM15k/8yaLUKn8EY2WJKXepCoyg+fbUpdZFFXznvwlxtM4YEFsDPTgm7jala6qf+REpNPKgISPOdgnLFQkgdJvIiQdt17A0KRSVhJEHNyBgEEDly5D5/DCRxF2mTRfATT7JkLeaiBvtZDRFMYAVwOuJIg4chQtQxkpK0EEK9dbQk94QkJcBpIhSJ8VSS7IM2pj4/zDc8EAlqRaBDTYyfopbMa0vrLlVmKVwmg5ZtsMhJGSArA+Rd30IGhJfhWFQtjzN5JluxIKiSaY8QX6JtivPmXjw8KPFvN5EMZgBrl05K/kl6tTvTtHY0nOlO0TUbSyIgotlVwwRSeXL4XHD6NFGEJvLX+MyJw6MaboxSihpHStbqayjST4cKt+qihrzDADNm7y/LPYf8oPslpN3tOV/ZR/6eNKLadc4jiInenjYBop5Q6mv+BGKVpO6yxLYMUBqebcOEUAuXJr19FFAAA4GbAQOuNxTFyRCykNsJ7ApqKP1MRKwDKOmMgy00QsW/b8NZCsUhFDXrjS35e2UaQlZ9RSPEnhUsm9Cm4UyYHmfkO8yUWcE93ttqK7vY/obj9AdLfrRccLlz1vEF1WdEF2iFgpUEU+5IwwrvYae+CX1LrFTO9WxfJROusthKRgCJLKZ35K5eSD9XEm7AHaUbkeOeC5ReLUXdkNRosZNLzGKCROTJ5OCb6VSwy0lG6joq8NuhER135nqlk1bNwekaw+lMAvqGRTeFofenhBqI8pLs6TkmUs8cqp8llUSC2XbJeyqOkk+IprujflO74wW4oTZzF7MD9xeMMqlE0jYRYkGTlfZQvgxZDzuqkr7fo8LQrnWbZe+Tp/yDAHPxO2AwiL+3Orw/nqfnVleG5+ma2JoTtAjoeRLsc0UV6YkCI6yzWINDHHsL6iDpl883MZuWrgC9rpkmMUCC65x8XbqnD08/fcgKQL4F8kJNdsb8p+dHzXMNdaJsJaJpSaYTasdmWwbXEknvBjFTfw7WX78jYGw/l8GW4S0qiO52uq8D56xjYfzcSNOpsKZTbaW9TImhquY3FoDSOn7G1DEGZHs7sDjPfn6g5KIaxLO2DwKZUeWNpvkZKoHooegtgRf72+ox6oHHIXIhtPBHMENZD1TaaypuI2g79dD8+7sWzoN3Bb9ZRA7xzipqU7Yf6EWwUMNv1aCC49Cg4C4FPD+TSV11le6ZJA46FkDfMOwPgLnb8x9o3PPsuSC+hpVrKLEQRlbLnHCved3LydpKOSDuofmTHlXjv+iDwOFn4s6t5DjSpuESW6BfZNVajc1MDBr5mdQ8ESg0iGLdRqyWCT1TfrEUQ+jg4FQw2ni3PzsMqrLJt8MeRiKrKTea9v5B2apsSdSgtlBlTKBUKrrMcvmHhWXodgb10VTx2vlO5W0Z3DVBza9QhKbrm8xHHD2Mmv2pD6gWWx2j7EiaL+zxqVxS6qepj5NIX6Wqity/XaiuJcyrq6XKer0styg54u12lpoY4utTq6LNYxFBJqqF5KlXXIWvXcWl0k83WpKuVynVLuFOZQNFjLfzRdeGPq+eQX7pRYgzWAURwGV6RgMr+I8mNnjsDR7wsnJBuhvw2oqVWa4QRv6Q8dYnfWDmy8sLo1J2wiPd0qVkeTrKbp/TbBpKlfizponRl/+Quaqbg2YQ2G0K8WoRQNqbpx1CocR61P4+incXTNOHr08cbRne0GT6tg8LQ+DZ5/6sHz6CMMnvS3GAzLTlQhsRSnKPvkJnmW1ysJCTig64IcFT6VnbBQFGnLKN0zSoZLTjZHxsTIXW6SVo0IppJgSxNOcFoLzzu4LdrmQP1k4UQiGnygAZhkIMNTPqyBZoe9iHVbgCULBNc54/E1iM8f8euR2JCzGH92jtIZQO7I3e2RNEQJ0W0BS3MajhjLXG+FJRvnosU0FmLeHy2ykkwOMyJZ/B8a6FbhFIPzJofEKb3BPcM3EtJKjqKgtmM6P28qd5aCJ8dvMuljP3/EnKV6PV/wdZNLKQxl1tjczfOgXuAupotIhqenP6U49E3EU8uXvX6Z0BvktrsGoasc4yUi7nG+c8v0s3PIkEBVM+eTNBRM5+2DIR1lk0Vyd6lRa7H//eNVBFqeKEJ41arI7cNV5HYrFdlor6o6IiKsVZJ8CR+mJBLin0dJknPW1PhozdCGOWucnSxoebZ2gjI7EhPj+cI5QdkSLAxd+u5jcDjdwmlA17suVeSzFD2fUhXm5dQZFZbx/elTyhIam/FLJvgy0vQAz3Q2siQcjFoaaEHTmcgtYPFo4Gf0ZGDqzSdnAxdAx6HjR7jY/GXoTVgAIA7m0AeMi+jDwHUSBnMSxrflkjdzJqQeEtR+z5/gES90zQ0IyS1VtiBQrycHkNbfB8EMCVhbIuJpY3V6EHIElgnFbM2XpUTcrDpkUb8bOdNR+doJy0q+aDV/cSdOE6/my2RXoEgprYntSDHwAlrJ8bR0IwHKDMyDUqIsDH/rShLBJTnRc5JRODpWS/nJZ1Xo/Izf7yGvtNRf3OU6/hU9wLHRJjPsZ6G8SYEL6L0J5gK55VbkcqsYptPvnSGZios7+KwS/X6ibH3gYwvOv8+gXVAYGlKjU+dTfJXmz+mXRIPw8KvX2JBQ+nMhniZCsZOSH9HTthHObFhQAuHMZC0W6jA7+kjoikpUNHhsm0WnZctmo4Oda74SK6hxLEFuEpVkO8rqj18zRMuj7Kpe32vqJCG2iOTIZq3IpqI6FjdC7rUI7U+tL15TjJBlml+Cvn1x56H+UfXT4/H+RCjtJk5U0zWLUGqYwyjoY9bJJ04ETdCxgrDHeAI7ykzoEQytonMtTz3a3JAlp6ctjK0uEhNXf/haqn+ySkm9ZhjEQB67zXrfdMmkpKWt6Ze5RoXY7WuzCaXxY2uN02lq3Xi45t9bg/OG9VoV1iinLJMttHNHtqBBVb9/sEpk6BtblwC6hQrloIdZ718MlGjBOhjo+5O5oRLWdzGkpLqpUtYDqpWoxJJWKqostv6M8kO7CA0FYTBmMHcUaF8oRw13WuEnSntVKVoBSAMvj4IliYot+A/yEbIMGqOpE0Xfe1HcAIsFLF1/HKAo+TUMqeV0n/AQB2OnHK5dJsq9ILaq+DVnSlF8ZSZBYH2bYuO1JCXhFGclL06j/O7f/C/usjayOn2nLODBCzByrBVlSm/OEFolQ+ZRVKPELtMoKanJUpy8lDgAWyNWnOwsZXFbCkAUk7k4t6EXBRMnQ9kktdKZUVLkxPTmgXJiyKmc6L0oJSWxWEwcoEhMSXKhmDiAKqb0c97DQnu1ZY4u8HSViA8Q2QixTr6MaF52IwycQspj6G/wohJ6XQW5jiWDPScLAGhAUSckZp+y0IeiW8WAO/JMdfGyZX7EhEpJat5CXKgg/q9WOfQ9T/HgBuyICHSXiW7UdEJR29VDkPOdLkZxpkDhI3a5nGLW/yRfGnQhXSm1P2ZsQ1wpt3cdj/qGrhY38uOhNo023VNv4SE04gyBXIKXcz6MS/yneW/DekBJiB2nTDzVqPTccHG/DJpT+Y5gy24gQ9Z2BEIya+klS5NCG3nJ1iTx5l1y/Nt8YlIqzYYqkQSJj+M49IYwmJdLtIZrYtUqgbhxgDdtfbwAH6eoHbkVkPVOfQKV9Oa/BXh5Vl0lwpMzAejTS1Xqz9OdyeOgko9jpNc0FTawvBG2lUASwkwiP0BHheVMvpYGKg/KQP7xeMDOKM8DnoelAN0vXJd3PxRCW0fr8pS4BZZe4vQR49AJzSITMwWQw3PCMevLF/ReNxg4lJ4OA0C4cIrG+rG+YBx6hFd6ef7k8dQDZl5Jm4KFKR8f3KJXibi+uEsocVelnpKmsRcQ1DvtjFASAxBs/oKVIKmJm4TpBBx5ioeDNKiPoDhQKdcaeDHWqneRUNriRm0sBbVzuQwVUhuEyaZwUvxE+44y0bEvyuFbm9hV7HyZexaoKNMgRcWgSRh8L9fNRhvtL6thk5lySMSHlNCQ9YTHb3PKkgmBawtunBpdiOfsadaiPJi0KQdO1FUlkp6PF9Ok8AntdZqc+beb9JgTTRFwOc9GeQrcyLumxvF+dnNHAkWLXEsyo2/ShRxxMNdhwecMCV4kHCrX/Wz3bIpFEzI8+iphJrdK5lFZSobL3gXkleZ+lnXTDoIe0NBv3ZBEInZH6yYccuEudiEmjcClGeDUikAfXqX4VgF9oeuUJSh2nSGfDdpELJ1BEKnxII5Ajn3J6K2UXek4MAE1NF/okqEoW9fCtHYxh2GGYPqzMJi95papfhuhtDoATztxT5bs+Aw+H06/lm883w1uGi65Bg+D5kqBpMP5C2DwtitpcY+ajWXsiVnh66bs5Pt/5inAsXu5iOIZiwEqGRVSLb5xZxRFJ8s31O2jl5SEEdvPX1ZtaEkOIpZSdpmgpRYHsSX0r5VCVIwvhSs1DnMHnHzYlK7+CMx70XQWcVAqqCY6Z4l6KJZQXIT2tbYmixrVhglnaGM0Q+HgKtHCkqbQGH83Tjy6EPRWbTF8UQpxPfmgYa5WM0R/ganld+WQ0IAyvaX0i7sCxVu5c+yClLDIlLv1ufwNljc/NUYJAYwuqNNYS9GF0UjXH4ijlVJwIW6RPLCM84GHolxXOSeyOGghF3JN55YeIDd12TyzQJ6d5gSfE9D1Zd6Rai4ahcF0+kiyye7ooKvLAa8zpBzUjCG5cK49vMi1hEcsOX5cWql5iI40zee5Hwc/e+SmfKdBB5rTYHQFX3zigBJlBFfypve8PGfBIqKagTJVY2jJ4Ue4zhQkJy44xdAWK9yvNSN7eSv1aQmmZgd06o6yZskv66WB38AnwgmSihdXBFh0GHS6G/L8WjkmMV18lJZsm0ypoIvZY3PRkpW+wd2ki2bphuiCObRcHrFivWzOgm742JwFVoc8IaY/eVLKu1jAa47MqQy2UARV1EJwAkwjgOioaWgxETeZKEspM9Vkn6spAT3yN0mUQMHm33Xo9EixrDlASxf5ODJM46uvJJGJwFUFmLleAseyg6hIiy99KEnn3Sow3DtWPMftpo7lYiVFV3RH47xp2dRysFXW1aKsiworSWRVUFNv05r6JnHO11TV27SqUmixrr7RePMKb7Tdrq+rZL74Ayvr7YdUVr57uUddvb1HXWWT46XCLSBbDV8BswjU0WvLIWarAWYjIzROt2kMfZVsNXjgSJrg54fT/K4YBHj44HiejYos04F6our5df4E2PVc6E/2SAOg8shTeowJxN3HYLmijRvOQ1YLkzvuRPez3WBlyDdm6bcUyTuqNsmkYCNWDk29F0B/QshWunpzQchUp6tJ/+hMY8g3H9hyyTR23qKLwXArDfYl4wqx+fHyT8jYAcVRQ6LZXdQ5o72Cdy9k6Yy2cu/NWktflEkeENe1K+xy1wBGDYT7Hviebumyp+fTiYgSoYG6ETrP0ZHR2lS6+qHRUvVPyvWggF08zSZ3nI3MrvRalXcw6rb6FZbjwKhvLEh1U0GOigri+fcqSH1zQTSTmSKJ7X3g/MPDpl7yTuWdGsmlDXNEcd+KcT8M4Rp7ig5/rfR5ZUVCErQUx2XBZiGzX9Xpno+Zm86XzepLc1zt3Iki75rsswtcVtKcmG4VxboK1AYwcgfWskDVS5884df2fLzjauXThcSDRjWLX1nwggMlXzRwabhYBGUfB/LxTPkj9LY/RC93jJ58CF7+2LxcOmMhVleLFR+Mlz8XiV4owe5t4IM/njUBWoJLIvaFi5ZAc/ioq+yMTTcWPWAVaeEWHmvkzEtFUNnunbVgMxpxZHVeMhvdNpkVwgoWq+fj6R112XPOzSvot5rogdE4Hk/pZuDSNVPpQtj1u400EivcR6FCYqlSz7xkN9qlreK7m8LXRWVexJghraT5UtmysbNub1i6yUtqBGv1QwQpUg4RpmAB8o4y6fwGh4XcMeaY8nSa88W31PiiUidkReOIf1LLbotrl3Ng2s0vBgblsQmAq5lOZ1p0nReMRElhk5MGVmT2rpD+2sXbOegkuFJvr+M582Itax2ctPx9RHzxYJ181pK+WzpA2V1SB4SVqg9sBPjzqAQr779MK3is51+uEsnAL2tF0ZKltC7zZ/jzPb53/LIAxSbaz39idpzmOP193UeEXg12VhVqhcUXHjUBXgVB/EPgknKlcRFEMXibYz9qJAGo5ERGeBwc7IHl7M3jo4O9OCQkGsEQUA8Xfv2ChOToABe4G3SF1eHuOJi6eJvHuQ+Ud48OqJyODujEkoGWw+Hu6IKMrqAMu1qc8ziYTKaIukeRZPL0tI9zZzgEL3j36KtpPBCTS34Q08TS0aWzbFAhGMA8QEhk0Ef3M8AEhB8eX7bA5GlWCggb//j7v5sN0/jv/8Lfp//4+3/Se37/8ff/gL9nxnsSBvtNc0OmSXLyh0tJLCmFJkv44hJ3V07DfeKgiO75CBULlGNtOiTi6k+ellTqeaqOCrJyu8bu0atkhGcK0Wg0Er61ykA1hVc1qP/UG9GmsBeMYhLXI8BxZrtH6ZVZrDlRnWRvqJKDNCQyptcErFXZAQNrAA+vg8BPrtXgJOI5DzpgbytdJlIuxWQ2R3MF6ZAwBIGBYYgbTjITlt2K8O3rlz80qANdRoINvi9cpcfKXqpIfQVQQ5T0XhFDbmcN3Z0YjN3sQg3sKzY0wqQsSTVKRdpdUy+XUQD6c7eLayJ3943dJyzslh6sA14DBgvBTzYcugCg0jC+QcN1j33HpVzsNgQaaNytGbvCHQhI8fh/+U/KMr3uATnmPQEmKFcr0PL8+PTRNv8junDWA6De7ToW/LEwQxMemuA77WZ7feHL6d0ubnPAbAAUUmnXQF8ByrAAB7/SPQ7wFXqeXWielNQZPN8W0TLztExgxNQTs0xKTdlehSTxs3zkEi1WIrjscCYkDIYDksX+cXfmLNNnHtmFd+ZS7qbnYKcg6XHY6ZfkLmjk4rTfYVsQmvSSWqvZYq8ti752upDSwf/paw8eu3jdLnvt9wDYxHsB2V2+Jr43+/DLYlf3Wvihi5dPttkHGz7YZhM+9NkluS3cDIE52E0GYdq4HYL+Ypl0EL2PdxWaJv1g4aaJrkm/4oduG1nGnHoUo4mPdjN9R2L9VgqOOXaw0B16DW8TqLWQomWdnaESSdsFQERdprpFUXemHUmt8ZsGdxMcqiLYIHYxKrQW7FQGoSeY0PYESpMpXnZ3GFWu5EuqbJtygecN2Vi5bJpyLs2NmZxtLu2/+QaMBjIcS1ydrcQOmnfG8AhDb264z9kGil2zzhD4m/GELlnaNx6f/GQkhsZXk3igz4UxgAYAtXRhYGABjX0DJ+B2Dc9FAyw8T8eLltNrjp1+f2i7/ZbZGfat/rDvkE6zN+y6ZNjk7MFvpHfErgoCa2KGHsy5xO0+mE3lhmTe7V84UfmI2TsNneFHcfapjUhcPHChQofyKYYxP3foDwyHUDqa+x9jexo8+21sUE2SVha6WgXwqTOPMJGashrFUGh8rjVn/6+ZvWsFUazJf7CVTOVWPj01G3gFA/6fXVf7f/EjjBP0OqlDXsGVArF+WHcxNlujdt912mbXbHWGxDHNntvpW8NRv9t0+u2su/h/6CXaG1tMu0mbTDNpMjiOQ6OxsNGARGm76XxyFz+5i5/cxcS9ejOf3568d56Sy97o+e9B9aeTX69PLsOrp2+qE/L6qjr69s3kZOJdPX7SD95Uf6yOv/l1cjL/rv3bm94vz256e61JcHXS//2H3y57P17OouoP3z+7OGlOPOuyffP9T71q6zvv15M3we+jJ70Xz723e29efNs+aXnfXF62Fy+XP+1NZjHAH//66k01+O3HXvWnYHJxUv327fGTPbt5Mdr75YW/OIlvq8773stfL2/2nn4fPj4Zt4bOZfX38a2593PwfHbSfLS8etM6fnrZ67/466+/n/wYtn95X317892k+t23T6OTZueYvOn/+GQeVb2n3cuT2fPhy/dV5+Lx072TF1Dkq+NX4Xv76qTZ23v/e88+uXlsfvO+e/Ns8qLaOn7rn7x2bt33ez/247fV1snwx5O3k5vf3nR+efnzi70/yp39ILfXTtxei5vm1P/trPV/H+SzFnrAdp6avcabRhN/gydu63Htj+g7181Gv981O62e3et17U4P/dbUl9YkbXKtW2a324PfvV67b7ftfs7V1tD85Hl/8rzv43nbci72R8jE/uTefyR73SFNkxDbhrbWbYHd3msP4a/dGbackeMMzX+9vf7/y6P+J7sI/7tda3BP+61up9M1DHRR+91eu9vv0+dm07b70C1nviw6s61Or9PptToI0e70ezAqUWir026bTbMJTm56nwZSbzV7nabdpjDtZrMLIxx9brXafavXtBXqzV6z32k2EaIH7mrH7jHMdq/f6hiGSr0Do2ffMik3fdPsmDYrh900zWYLBgSJeq8HY5FpUeiW1bObVpfmBENW07ZgCFSom92W2bbbXUq92bNarSalDqNxvw0Sk6n3u+2ObfH8u71mF34o7z0bkXO893pmu9Vh0utSSTLJWC3gvWM3Zeqdptnvccm0TasJebXoc79pt9owTqpy73SaXSBKeYdxvWeZFLfbbPdQBjJ1s982oXxtCg3l7PU7rKRQBhvMkpzc22bL7ti0fHavZXZ6bZqTCdm2eqbCuwUUkT6tVbByOj1Wq0273+63mjm5d0DuPa6F/W6rY3W4RpqQV7PdVXSm1ey3e22blQ5K2uG8g8RsKIdK3e6DOdVktWr22mBZWRSX21lKrXYtqAvT7LJLQNv9jmnSklp2zwTDqJ/TGW7PUYodphJM3/q9ttVVWpNlgi3VbFKKtg1a0u1TzJ7dNqHkrbMPiRl9sHfPzvJ67tKDhDOHtG+D+o8tc9x02yhIp0Paw55jdUdNq9MaUytn++C1CLwxdCUCbxw3zwY8CBGMwiCItWGIJFKAsuJAHFyJDew2UnGeY/JuDvXxNKCL6JL3xgg/0GzonWCDLCUkYBGMyC9efFGW0FOi49CZ8B36ysqkJ/z1GYfAMiTQ0uqWhDCmaw73SCtXOdyD3Tae7L5PKcMrXxP16PY5EE+w8dqdAbunPliEI/KEHmtVIMPP0fDZNaqGgs4PjRPFklJrjL0wSvKmJQOELDULy6zkmtAKeVMU5n8AodSbog==</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_3afdf582c2b64763a175ed6a40527896\")) .filter((elt) => !elt.dataset['step1']) )[0]; root.dataset['step1'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "class WeightStack(nnx.Module):\n",
    "  @nnx.vmap\n",
    "  def __init__(self, seed: jax.Array):\n",
    "    self.kernel = nnx.Param(random.uniform(random.key(seed), (2, 3)))\n",
    "    self.bias = nnx.Param(jnp.zeros((3,)))\n",
    "\n",
    "  @nnx.vmap(in_axes=0, out_axes=1)\n",
    "  def __call__(self, x: jax.Array):\n",
    "    assert self.kernel.ndim == 2, 'Batch dimensions not allowed'\n",
    "    assert x.ndim == 1, 'Batch dimensions not allowed'\n",
    "    return x @ self.kernel + self.bias\n",
    "\n",
    "weights = WeightStack(jnp.arange(10))\n",
    "\n",
    "x = jax.random.normal(random.key(1), (10, 2))\n",
    "y = weights(x)\n",
    "\n",
    "print(f'{y.shape = }')\n",
    "nnx.display(weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13b52d61",
   "metadata": {},
   "source": [
    "The rest of the guide will focus on transforming individual functions. But do note that all examples can be written in this method style."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0251e7db",
   "metadata": {},
   "source": [
    "## State propagation\n",
    "\n",
    "So far our functions have been stateless. However, the real power of Flax NNX transforms comes when you have stateful functions, because one of their main features is to propagate state changes to preserve reference semantics. Let's update the previous example by adding\n",
    "a `count` attribute to `Weights` and incrementing it in the new `stateful_vector_dot` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a4fbadb3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Count( # 10 (40 B)\n",
       "  value=Array([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10], dtype=int32)\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Count(nnx.Variable): pass\n",
    "\n",
    "class Weights(nnx.Module):\n",
    "  def __init__(self, kernel: jax.Array, bias: jax.Array, count: jax.Array):\n",
    "    self.kernel, self.bias = nnx.Param(kernel), nnx.Param(bias)\n",
    "    self.count = Count(count)\n",
    "\n",
    "weights = Weights(\n",
    "  kernel=random.uniform(random.key(0), (10, 2, 3)),\n",
    "  bias=jnp.zeros((10, 3)),\n",
    "  count=jnp.arange(10),\n",
    ")\n",
    "x = jax.random.normal(random.key(1), (10, 2))\n",
    "\n",
    "def stateful_vector_dot(weights: Weights, x: jax.Array):\n",
    "  assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\n",
    "  assert x.ndim == 1, 'Batch dimensions not allowed'\n",
    "  weights.count += 1\n",
    "  return x @ weights.kernel + weights.bias\n",
    "\n",
    "\n",
    "y = nnx.vmap(stateful_vector_dot, in_axes=0, out_axes=1)(weights, x)\n",
    "\n",
    "weights.count"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "322312ee",
   "metadata": {},
   "source": [
    "After running `stateful_vector_dot` once, you verified that the `count` attribute was correctly updated. Because `Weights` was vectorized, `count` was initialized as an `arange(10)`, and all of its elements were incremented by `1` inside the transformation. The most important part is that updates were propagated to the original `Weights` object outside the transformation. Nice!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7294661f",
   "metadata": {},
   "source": [
    "### Graph updates propagation\n",
    "\n",
    "JAX transforms see inputs as pytrees of `jax.Array`s, and Flax NNX sees inputs as  pytrees of `jax.Array`s and Python references, where references form a graph. Flax NNX's state propagation machinery can track arbitrary updates to the objects as long as they're local to the inputs (updates to globals inside transforms are not supported).\n",
    "\n",
    "This means that you can modify graph structure as needed, including updating existing attributes, adding/deleting attributes, swapping attributes, sharing (new) references between objects, sharing `nnx.Variable`s between objects, etc. Sky is the limit!\n",
    "\n",
    "The following example demonstrates performing some arbitrary updates to the `Weights` object inside `nnx.vmap`, and verifying that the updates are correctly propagated to the original `Weights` object outside the transformation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "76c58a29",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<script> (()=>{ if (customElements.get('treescope-container') === undefined) { class TreescopeContainer extends HTMLElement { constructor() { super(); this.attachShadow({mode: \"open\"}); this.defns = {}; this.state = {}; } } customElements.define(\"treescope-container\", TreescopeContainer); } if (customElements.get('treescope-run-here') === undefined) { class RunHere extends HTMLElement { constructor() { super() } connectedCallback() { const run = child => { const fn = new Function(child.textContent); child.textContent = \"\"; fn.call(this); this.remove(); }; const child = this.querySelector(\"script\"); if (child) { run(child); } else { new MutationObserver(()=>{ run(this.querySelector(\"script\")); }).observe(this, {childList: true}); } } } customElements.define(\"treescope-run-here\", RunHere); } })(); </script> <treescope-container class=\"treescope_out_68b34aeff02e43029cf4dd2722226fb5\" style=\"display:block\"></treescope-container> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_68b34aeff02e43029cf4dd2722226fb5\")) .filter((elt) => !elt.dataset.setup) )[0]; root.dataset.setup = 1; const msg = document.createElement(\"span\"); msg.style = \"color: #cccccc; font-family: monospace;\"; msg.textContent = \"(Loading...)\"; root.state.loadingMsg = msg; root.shadowRoot.appendChild(msg); root.state.chain = new Promise((resolve, reject) => { const observer = new IntersectionObserver((entries) => { for (const entry of entries) { if (entry.isIntersecting) { resolve(); observer.disconnect(); return; } } }, {rootMargin: \"1000px\"}); window.setTimeout(() => { observer.observe(root); }, 0); }); root.state.deferring = false; const _insertNode = (node) => { for (let oldScript of node.querySelectorAll(\"script\")) { let newScript = document.createElement(\"script\"); newScript.type = oldScript.type; newScript.textContent = oldScript.textContent; oldScript.parentNode.replaceChild(newScript, oldScript); } if (root.state.loadingMsg) { root.state.loadingMsg.remove(); root.state.loadingMsg = null; } root.shadowRoot.appendChild(node); }; root.defns.insertContent = ((contentNode, compressed) => { if (compressed) { root.state.deferring = true; } if (root.state.deferring) { root.state.chain = (async () => { await root.state.chain; if (compressed) { const encoded = contentNode.textContent; const blob = new Blob([ Uint8Array.from(atob(encoded), (m) => m.codePointAt(0)) ]); const reader = blob.stream().pipeThrough( new DecompressionStream(\"deflate\") ).pipeThrough( new TextDecoderStream(\"utf-8\") ).getReader(); const parts = []; while (true) { const step = await reader.read(); if (step.done) { break; } parts.push(step.value); } const tpl = document.createElement('template'); tpl.innerHTML = parts.join(\"\"); _insertNode(tpl.content); } else { _insertNode(contentNode.content); } })(); } else { _insertNode(contentNode.content); } }); </script></treescope-run-here><div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtXI1X2koW/1fmpWdXWCUmIeHL6tlgAW2rrWJr27c97CSZfEhIYjKA+I7/+95JwkcgUO1Tu+0pnlNw5s6d+z2/O8G+jOjEJQc8DQmJdD8gvdD3KfoLBX7kUMf3GigkLqbOiOwh0/doycQDx5000MD3/CjAOoyPbYeSUvxLAwUhjLhOREsx6xKdBDDq+R4Ma1jvW6E/9IyS7rt+2EiW7qH0N80FAuDnGNRuINOhQOZR4tE9NHC8UjouCsI/gJd/U4qcW8ezYJ0fGiQswdAeCrBhwGDJJSZtIEm3mTQeKdnEsWwYEXmF7edR7IByM/7ph9LIiRzNcR0KKuIh9We0JcejoeNFjs62Jclsqtfdy93Eji9ndiyFQw/2DGEs0kMnoIgZYn8LB4Hr6JiZdtfXKWFmCgkebB0UCsX9A7A87BdRZBDTi9A+orYT8Rah5+CWU98ghSJv+xHl43lQjVDUC4jHVFZ1xpUt+vNr3swR9gyXwLQ3dN29ZAcexOz6vgejhbEf9otoUQb/EobYVGaYOjobDEho+uEAezrhPX9cKMaBABsUVmZQKVn0EpWlIvBxTFRYkpp3iWdRG+3vI4GRbBQ9JHQYemB3RNyIzAWzhx6TbJl1ZDsmZfLFBOzDHfys2aEA4ecZ/pgPyfWQRFT1nEHsrnaIB6SQ2KTIeOytbBQMIzsx416OjtMt9hM1Nmh5fxmYFIkjqW9ZbpK+vTjFIFoDxouNEJfuIDKCAE89yaSLf+f7ZMKMzoUcEygl5nUXR9FbyOKUb4Gb8ewNIAy56eZ3RbAnhH8c4wcvd/MSwHBGKGa4z2XrDIco1kBTcrPPCRykbkhXSXwPRARjeDC1KRnyLVBga6a6c5CMSb3TfEr9ASsMDc+nBd70XQNrsNoDtg0bR4UDF2vEPcjO9JI94jUN3SZ6nxjFIvoXM9208FA/aCCBFxUyWC09bOy2FKsMv+/lltk7Pi6HPaxpIRnF0R1XxxeVmoQFYU6g+4MBqLVAgeMXc80SCU5Etv0RCYs59Ck5WBrC11jgKMDLNOcU3nCgkXCRoF4TlcqcIGJV0lqUqSwqojIjIEaPGWVOkBb91bNhhMNCqaS5vt5Phop700of21gMblDku46xiTJx9LeI71gIkZAJF4A4xAanZ7TEVfaTHoFxDDWQQzGUcrY4EyNrDk9wSV4oAbXhRLDpZHpILhOiAxSHYqOhESirC5Z7ocev/ChKzr+SyA7A9OCE0Jnt5XhxaMZm2LBnHDCrOxs47EcEW5Cq3urqR0qpmQxsaf6iKX1Gwvggb6Ct/0iKpm/9SPGyi9YKWXkGIZkf2cbDMGIODHyAMyTM2deJHm/bOBXijUpxAY7Wxfjj7DpXj5IburoL70Q90wkj2vO9pAitptamVOIlhWVTrqvQ3xY/8fiyiEyrAQ4tAKCJGHFC3/3N3aAUBxNtCKXRyy1A8+m8oOUQt0QFhgT4n0/8HyLKBrfUPWydYIgKB7uoOxlovhuhd0PK9DXQYbIS3oMJJEZpTLQ+dAJJ5R3A4W7HmB97FJY7OCLGrH94QQT2s7ca5snqGLcLfJ2dwlktk/zI0SK/3M1X8mMc9XQ4B8Cws/XYpJnTY1qnN+25tCa75aLp0wPMwBSXsAeOjZFhcXGYbcLgboi9aTTHbJEYIQIWAwBS8of0YarMJADHOMT4IytJvCX6wxkEfkixt8JbC/0+8XpsZF6Mvm3dhWUL9py6+Y63dbsXN549lh0LKCxJFoFPMnaBLmRYbIEwTLDZjJJhM1DV6OnQzBgh8VLls+0iSJklfDQYOUvG9GjOJL+OXb0APS00VABoYgzORxSz9TN5n0ySFFElkhiAqyBMGKxacMf1ELsetCi9ICSmcwNMMolXixMP+jHMkNYYhx54oDc9KqbeNU2si+UcwgAamhXHxfUwNVI6VEqdOYeUjRiI47BkhdhwwG0FJJYVg1g7yIcksQgSQLyKbu8kSQOdCCtB8RBKzbwiy0qtfpyDAK2U/NIsOJd9vS42l+lm0DvxIWR3luTZo1fgZd0u3kfW5wnnpDWIpcptEDLNXawH+MS1wVsFEdRZXJLyjYCKFFmTntpgvrR0/7XTfadtTLxyxnI6WkqHn9Ga80YstmCJ5dUwShUUyQC0W69YDoZ6MIL5Voj/lekVE/+yanWPXnEz8R0PHJ/k5mDu5L9zeTC24cBkdwd5yI4BbTBeLk3SEi5QgazrGCUcwCQuDgCBfbuLfbiD1+8wFzQhIjcAVYy1NI8hR94WGVPMr9zyTJG90+Kzd2loA4dlVb3M7V4+4X02W8vnES+02R0j+kMNQzzhzdAfFAxfH7LbKJ6BgogfYXdIINCKfOQPSCGGCux2kr3zSaPAbibv2SpwW3BwFmd3wZFNCGUXxmSMDrvdLtOmy8bY9W88yYckvuvpTjy98N9/p+2JTqag5eGtyuLlkMcuwN10bJzmsMxu76JQb6Bh6BYYbm6w+d2xb5rSngbIvCLvGEK9c2KpTTV+HZ+pqh9/ap6P4d+jtqq21E2v5kBVrb7/xjhuNQ/Hn1X14vPha/XkuHmotq2b46O3No2aJw6xyu1Xn6S3x5XPo24wdN6fKBfi60/H5x9PRpcnt/T9pN0+3L60+hdO85VgO6/Ohq9bRudKONJ2zdGxEVy/qdjXl45zNjzxOvaR+YGqHyrN01BW28dev1XRPwyH3va5cq1H/fHIbLu71zdWy69Z2utxpyYeqbueeq68DcPX4vm2dSucG4L62hSt0+rhuHMlWYI/GZ5Xq4OWWBkffaq/s6yAXPQnMjnWbhVdC991KFats+PT8SscTaKz4fHxp8tWe6y+PwuOPxsfdne3repF9VOZCuab99fqSAGeb9XTqnoyVgfW7Xl3e/ilS1qfbiSzot+eyudHE2XYVN/cNq+CdlB2js4OW8KX4Xu5W/XM5tvWUftkoDrbtVFLsj3Rrm5rH8efrsZH4ehV58Ohd2W2Whbdfqd/cd2qUj98PW7W7Lp8ctLpljtfVGtwrFw1z+r0okOO6q1m87hTfmXJ57uf9YmmdsCnH9/sqmcdrJKTQ1c9um29s75Qq9J8b717d/yq2XfOFNJufjpstnVHCOzQDzyIjeBL65V4K/a75qFJ7ckb78jA7ejIFE4HndZppWmo1x8/BphG3S8Dw8BOXTJv6/IH5+q6EgzCyjv/82HXCTuD0etOuXvZLbdbkt48My+2j1w/6MjtaKxg67pSc76Q7qkbXHrNo2NinIRkeHndORyIl+2w3+3eKFLl8jIaqyBREcUPf2hhKw7rLXZk/hf+mWU/NvwAeod5SsaPrHie30Cxk+TsV+C1+SGAHT9DiRvGpJcF3hAeno4KSUuZfcIFKXjhs/QFsrTlZGMRlAfGgvXQrPHEY+xQ5OGRY2HqhzxwDjQfhwY/Dh1KLsgNLcx5MUSR8Jo/RoEjvsAtNNjsAQrscuEMCHTihekTtpV1IRlAt7yy9G4HSYIgxEgKii+g1kJ8VZS/70IXzc2FY5dk0wrGnjlx6AVqY8eFwkZ9xIj/iCsboE0P+jqoxg7YjGCDXQJsL9oufRj0jcdA7EIBxeVxn8uAqgby+65uA6Su1hSGfoQyEqtVXqpXy5WyVBMlBe0CAAJ18/Aku9jmUvbpY6bsZfZynwbECQB46XjBMD3IuPjI1/wbLpdJig5gMkEGoGO8OLtv5glHChNR9o5iSdLs0c8d/NOlzI5AsZkuM7nUdsNsyiF5u4yPn6gwG021X+E/RTlcjlbTR0ncAYTJIQjCMKmACrKAmsUd9B6HGKBrBUakdOjCp9htoCobqrGhrFSLW2xNoeZWdpgN9dyRmx3mVtrgLPbnNlCvkOpMlf310fN3giU/SJ7D5bGDCnme3prfrm0h3ztkNWV/64FFNX6qWtxCs6u/fY6PDcnNE3w2BxAqmWS3ucv3ljAZlxA4omz4nGp5sD5UHhipsxB9tOBbuSDkNs1zBzHM3b+HZx1jn5s9hixj2RB0XBfNqizXDai6FQ3LuGwIuF4zldrSpnmPL9PIytf0Ct/wMTRHeRaZ9Ycx4ZTE8WhZKojCTvGf1jxkv+GpTKeRq6okKlqlXDVlXMNy3TQ03TD0mqCVdZHoulzdVNsfJzvzCuL/jUW/nQI7D02TTZXlaetD0vRtrBIJyePUiuV6AAXhoxNB6XRuY47IdgzIVgA2aN4Nszb5u6L6FzHg/d/WBPgmF6w77Ba/8sIdQDNwD3eu5kxx80G7Qr+zSj+7cYJweWBe/ZRH6/pA34yb+iT0iJs517ZWn5Fs5XL+KQGW6UKB97wbSK/QYZK4jsZnzReD4OdFXYkb1sVGMvsjcNe8EfjJgJeombqkKJJcrVRko17R6oYpi1LFrCgY66LwI4CX6fo4BQpI2kHlx8VfZbOmVQyzZgqaIUs1GRNZkCTJECQimXJV/IXx1wMM+6vAsKQkbIYRizS/gdgPNOFvKPadUIzL+45JfDCdk4Bg9tWVYEJtCDxfu0KYIuFGVISyRGQi5Jn05zy+vxfbeWTcCxiW+Q3vfjC8m3liXYjMCH6DvAeAPLmi1GtSuUaqVUWu1ir1OtaVWlWuy7W6LGnKrwfyFEUxq7okijWzLlfqZU0TTLGuyEQ3y0ZV/g3yfimQN6sKm0HKEtlvqPfDDfkb8P3igO85DvTvhX3sa1iAh/yAhHTy1I9E/3xWHJVRbZ3pM0SPUgyfGPzkJGjy55eQby9upOoeTt5+rtvujBv+FL7ey11A9xQJcw+LJ38Ryx1IP7OVxXtaWfxRVk7/Mpk7aGM3Ik8A1J7J0NI9DS19fVIQ8WTooYyIS9iC6HtxxNef1rlPfK78jKBw/bfj5v8DA1r64zzuwVDyARcaU0rDGR38DwJA6Mg=</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_68b34aeff02e43029cf4dd2722226fb5\")) .filter((elt) => !elt.dataset['step0']) )[0]; root.dataset['step0'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtfe162zay8H9fBet+SKolmaS+Ldv7OE7Spm1Sb9JsN+vjx6FISKYjkSpJ2XK8/v/ufbzvBby3cC5lr+TMACAJgKAkO+nps0/tNjYJzAwGMwNgMPjgvudfGXFyMyUH254fz6fOzZ4RhAHZNnzvYHscRuceGZMoIt65bXVG3VZv3Hb6Tnsw9kau57l9c9RyLeK67d724X48dwL4jfQOm04UOTdX/sdzNwwSxw9IZNwa1xd+QhoA5xIsKJo506FxZ+iAm34wDgFlDCmNsTPzp8DbLAxCij003HAaRnvGlw79GRozJ5r4QWMUJkk42zPMpt0hs6Fc4jwiq4vzg/kiOU1u5iCQyAkmZPsMWLgiUeK7zrThTP1JAFz4njcFSmN/mhDgYQLUYsgnVatmhFCUn9xUzWandu/C9i7CKyqoIun70QsWsxGJgGAQJtW9cegu4hqQHYWRR6JG5Hj+It4zWvPlp5Fkz5RpJJ/qpEd/hry4PcOaL404nPpenrWi1GYMkCSKVXtZpT3KQuLPAUcy5KExD2M/8UNQmzMCHhYJpI0c98MkCheB1+As04J0DI+mAAtUHM/zgwmzK/cCyfoBaKhBrkiQxGlh176XXOyB9pIGMgdZQwM5G0/D6z3jyo/9ERpOsVofG37gkSWUbJrm6lqOwuWGtQyXjfjC8bBok/6H1aIVqvMEGxJ41fUVyvgarGDLnfruB89JnPtobBo6KNHzGYljZ0IE60lb9N3+LutL9pOIkNgN56QRLYLGBYkgLXYjf54Y1DYrznwOPDgogd3QTUjSiAHHmVUOt/AHio0TI+XCODCq1ZpxcGjcbhkG/D9eBC6iGh6JSeRDK/9I3oI0+lVsNABgGBFJFlFg0NQjpNMcR+Gs6iThCIDqRnVGCc6abuiRExTlUVI1a7UhYN9tlRfzHMSQtOy8IMbq6CYhMfD5oPJSImOkjVQCcm3wgiitKiXfHC3G0LdzFF5BhrOO6xfB/wrPtJh7ccxYnpLEOEZjmjnz1989eQqWOVRrMyHJMRijHyzCRUyBq1fOdEHqzAwBE9HSGiLFkROTc9oa6kY4HsckYXz4Y4OhGvsHhpliGAI8VMcc8lSGmafcGWQaE4HI4YFhlRAROWtOSTBJLoyGYRdIW02ZeEqMidhN4owiK/Jbo6onbdWGOj5eOslFE+QOMsuI1Qpc5OV8bVicH0HTkVKh07yIs1PzDJmygAVGrmbscPJGGZKxY1gcUdQOK2yyqjDroYVZ+sJGqwqzH1qYrRbG7f80qhuTujE60zfamwBcJ/cocmM/uHhNgHqVl/eB3NA+/2/c7Kf+/KUDth051y/9gP3Fd07iO2eemmVGPU4cGM/eoH/i8SKqgJE4i9yC0bK/8OPnfgADQ5Vm/fOfzIRgqKoua8YuIhj7hkUa7Rwvq+AytSzFmjMASisGzwyJfUuJfZvB4A8FmIaTarHUHY79W5SAUvjbPLyuLhlA3bBrtcy27wQrztq+JEfjQOoDMJ/JU8lQxI8Ww+RfWlkVnpHVcaZCZrWWMtJelGo5hZo5y2qqd85QbVhS0/0MIuPyD/rDmh1oC9uMcbClaJ4s59XMBGQZQDPMzRtsJwNLq5apPisDRJWWtMsasiyWXUVwbPAAwUJfC7i6JpM1PWwdQ8F0UizFdDJi0GIsUzaBrHvFpj4n2MuitcO4u6qx7nKaKS3eslIi2tHkk7WVG30qyfxVZ6oqIZSrR4Jw5gfgY0SZDftBVTABXbWVro+LgLIg9Hb1NVTShoLoktokphTdyQxrFciFv5HaBHL5EFsYBZLwDTrSb5IIPG421su+WzaIpEOXNMS8jyaj6le30Z3x1e0Ef43uau+1ww3695ETg1s1eViJAgRORwKY1dwARNO27C60zwi66GbP6tjwPMFns2fj8yjvpHK0Q8Oy+7nseWUqdAJU0Zp0CkInMRVJoEpFYX4+PXESmDAF4AOBPuD/mzpME2lSNlCCkqtopj51+ODPfgrCnSxI29mpKR5aFF4DPAc89c9SA8nIXTJyl0AOYDNSlyIpPviE16eXZ2IqFJIsm8j/a+ImVfQuLoF3+OPXDasu+Hy5Rd4VTIsx6vkTP6HO80nkz5wIVXVKYStfjulPpQ6P1rjXG7Xp43jcG5uEPtquY9ouffS6ds/u08dBu9sbefSx73a67VGlzgmSVq/n2jRn5I48mz1avRFxxxWAoWJS+XpDIMWTOeuN8T+K7RC3R/qcs9Gox3noe+O+w1MH/UGXPrqdkel12GN74A7aGWfj3qjrMXa8kTfqM/YHxHNIJ+NsK+POJdPpG5hFAUu9IctQJi0wMxn7k8KcxYMe5+eAHAN+2sNR2wMl02lL3eBTGD+Gns338rkMI1jPOvTUILhzRqEFK+QGQpswcFkJP0zdi2qn8zUGDWqV4ZbGkKAoaIg9yg17wP9rw9U0e6ZCs9CwOF10gjPa7OU0s9NTs27k/5/VpQyLplrFjM+CcVYrzONkoTcxQgHDAFTZzSacFbXJu9ztl7yXAqkUShpt/PiV84pNHmtiQy9IXOz47qs+DTkcFOxOh8oD/tYEyg9SIojYEkRsnYnjb6laMsVYBVWuxior66ymuOD51Bx0+AL0ABOZG0nQTIOoKpCLJrAA8/G6UdBdUaS6YfpzqWvlyPzpmmMitu6vnodlactarbnGg1Rn/ilUl0tUp4Sy5rO6Pa5uqtaaRqeqiA9cHk58+RjGu7yh7O+kcIcYBVP8noe20/uq+74Kf7DKH6z0lQ1sXaYltD/rXpjlmWei15m33kyd+4Z5b3Wafz51mqvkbj5Ynfcmq1GnqLrcbUkVXJN1m83x7+97aAsskJDy8Ec2AY01pZwKJiWp7u7B9qLMstY5kkZl7sBkIeHR4Uoh/PZJHqXcqtLsU9rXnoluJwcBfhbABgytxHv0Pj+/96kOhCVxE52MsrBNncZt6kLg5v4q2Ng46dx7FJbaJU5iT5woiZ/cPEXQbGJO5SJOuFB+vbNiCqjNrhstfQ4ItrMWossh2vi3BQa0FrLDISlG+yEYLfzb+TRM+NtFe8lnn2KM2Q+ucB4PAh07oCWxsTIv+GvDMr5Q4pG595RhJ9GCrLHCgEycxL8i2RLifr7CmcLMnAm42wtPWoJQfTYhXoOrvhlOE/tj2hXXmvF86ifVSkVx9RhSuli5XzAsnqNzGhD0fI6wUKaKdyoRPhP7eJTyPKI7JM4jMidOEp+HY1ytXkyn0jiuCfxJZIfGzo6vjnm8hccJcFM3Yt8jnAfOJWNZCAgWZAiAr+g2HS4dgK3JwJnkaNddrIzKU7Fb0QTXmMgUtozNRFUyXq8qlwcb71MqBVWK1foHtC8seAf5CMKVk48iBW2Jz6UeguIMpe2pttprKY6iBVepbPAT+vVTpXqnet8vn+RpMzQO4VlJdTUKFuuv7XhEt0gMdjHgR0/jc3oa2r0iG4thc/GWLKWgG8GW4o4C70Xg+S6Jq2ok22fp+BCDnVBUZQNS2qGf8i6BdsLpWhIgGdAT6LBTy4Kc0wr6MZUz6snAqEMiZ1oRjY2W0Zwv4osUgTJa0caeiiRVjz1lPds0w6t5ylGdpR9Xzs7kgS8FPjA4VPzBn5/TfqiiLPUI7L7/6lYDfrcnJ5PAg8T3+tk4L3j/vuUyPB3VdThGQ1cW7kzC9arVfG+iD6Y8ddGssCBSZI8i3qNGlerMiT8QzwgXSa3yIDbPp2H4YTEvcJuu3xjffGN8wXH9SRBGOEGkveUK7ZTzVawOM9V4MYoT8NFo281MkPF2TteqK2fKdDHlVEYtmzkWOFwEH4LwOpDYK/EaBDyxsLJxaRPZYxPUiR69O32zbSKKvs0erm86Gc2HtICiXd1LbaxIhfVNtbZWZ6tbSIm+7somImJx+385rCiTinBKmiSKwqhaect4Efv+Ch9HtDu7+C4AVsBl6Afp3EPaYHoESn4zJ25hlfYcnD/n5m2Q+NO/sQ3fVY9g7I9uT64bDgVLZSeMX3R7uA+jzs3Po5hEV3TzDt8HS6KYULw0q1oFHzfySZxtZE71xdNPzbOmLyC+xuLBAM3i6MO3Yr90og+47f7AEPht/rYg0c0b8GfdJIyOptNqRd26LYqeVa4qLkukUyEyxQajFCabEIA0IzILr0i1pjPlooSanh9DJQL0PVRl1o3bu2yvMFQjTo4CmDggg88jZ0aETeAlxEP2IOovdWX0O7tHC3/qHfF95s/9ySJSlO/SaEla63WmIjN4vil1mUXRNO/JX+o8jUMWxMZAD76Ju13prvonTky67RxISCzAPmWhIgmUpomQdNx6CUOTSlnJEHHwBAIGDTy6DZ3DC4ki7DJtvgJoniZC3mggb7SQ8RTGAE8DrmSIOHIULUdxlZ0ggpfrL6EnPCERbgPJEaRkRZIL8pz62Lj+8EJwgCWplgENt/J+Cpsx1Ve+3UpUKYyWY3bMQBgpKQDrU9RDD4KVFHdRKIT9YC1ZdiqhlGiKmVzg3AT71WdsfFgE8WI+D6ME3CCPjvy14nZ1anfn6CzJhbJzIopV1kShZZILpzjJ5VvhMcFdRBH01nJiTOZ0EmOKsxgllJTt1c1NtpkGH27UpJqyxwwLYOMmLz+P/Wf8IKs76Xu2s5/yLyXeqfWUa5yEiTM9DqexUu9w+iselKL1tM7yDFYdkGphGqcIoFBv7T66GAAAJwcWQmc8jok7ciGnCd4T+FT0kbpYKVjOERNZ7pqIdcufvwWSO1TEUBbu9A+kYxRZzRm1DE8yuExyr8NrRXJgud8Tf3KRFER3s6nobu4juptPEN3NatHxyuXPa0SXV12QHSLWSkyRDzkuxtXeYA/8M/VusdDbu3L5KJ31BkJSMARJFQs/pXIKwPs4E84Abalcuw7M3GJx6a7qhe5iBg2v6UbEScizKcG3aoWBVrJjVPS1SQ8i4t7v3DR3DBuPR6S7DyXwCyrZDJ7qQw8vCPWY4uI6KVkmEq+cKl9FhdxqxfYoi5pOgu+4pmdTfuQbs6U4cR6zB/cThzdUoewaCasg6cj5Ot8AL4acVy1daffnaVE4z7L3yvf5Q4EF+JlwHEDY3F/YHc5396s7wwvry2xPDD0BcjSKdSVmmfLGhAzRWa5ApJkFhvWKOmDyLa5lFNTAN7TTLccoENxyj5u3VeHo1++5A0k3wL9MSa443pT/6PiuY6n1XIT1XCh1w2xandpw0+pIPGHiDh7g283P5a0NhvP1MjwkpDEdP9Co8D52xg4fzcSDOusqZTY7G2hkhYYbWB2qYeSUva0JwmxpTneA8/5CPUEphHVpBwxzSqUHls5bZCR2DsQZgtgRf7u6ox6qHPIpRD6eCO4IWiDrm0xlT8VNDn+zGp53Y/nQb+Cx6imB3jnCQ0u3wvoJ9woYbJZaCi49ChMEwKeO82kmr7Oi0aWBxgPJG+YdgPEXun5j7BlffJFnl9DT7GQXIwjK2HKPHe5bhXU7yUYlG9Q/MmfKu3IClxyHiyARbe+hThX3iFLbAv9mR1Bu5uBgau7nULDUIZJhS61acthk8817BJGPwwPBUcPl4sI6rPIqy6ZYDbmaiuxk3htreYemKXGn0kKZAZVqidBqq/FLFp6V1xH4Wx/Kl47vlO5WsZ2DTBza/QhKaYWyxHHD2Cru2pD6gWW52T5kEkXnPytMFruonYN8TlNqr6XWulxtrSjOpWyry1W2Kr0s19jpcpWVltroUmujy3IbQyGhheqlVFuFrDXPjc1Fcl+XqlEuVxnlVmkJZYO1/EfThTenfkB+5ZMSa7gCME6i8AMpWcwvo3zszBE4/m3hRGQt9A8hdbUqM1zgrfyuQ+zWyoGNV1a354QtpGdHxRroktU1vd86mCz3W9EGrTPjL39BNxX3JqzAEPrVMpSyIVU3jlql46j1OI4+jqMrxtHDzzeObm02eFolg6f1OHj+qQfPw88weNLfYjAsv1GFJFKcohqQ6/RZ3q8kZOCArgty1PhSdspCWaQtp3TPKBluOVkfGRMjd4VFWjUimEmCbU04wWUtvO/gpuyYA50nCzcS0eADDcCkAxne8mENNSfsRaybEixZILjPGa+vQXz+iKmHYkPOY/z5PUpnALkld7eH0hAlRLcFLM1tOGIsc7UXlh6cixfTRIh5f7bISro4zIjk8X9ooBuFUwzOmxwSp/SG9wzfSEh3chQFrR3z+X1ThbsUfDl+k0sf+/lDNllqNIoVX7W4lMFQZo313TwP6oXeYrqIZXh6+1OGQ99EPLV++evXKb1h4bhrGHnKNV4i4i7nu7BNP7+HDAnsaNZ80oaC+bx9MKTDfLFI7i41Zi32v7+/iUDLE0UIr1oTuXm4idxsZCJr/VXVRkSElUZSrOHDjERC/PMYSXrPmhofrRvaMGeds5MHLc9WLlDmV2JiPF+4JyjfgoWhy8A7hgmnV7oM6PlXlZp8l6IfUKrCupy6osIKvj99SllCYyt+6QJfTppe4JmtRlaEi1ErQy1othK5ASxeDfyc3gxMZ/Pp3cAl0EnkBDFuNv858icsAJCEc+gDxmX0YeA6icI5iZKbasWfORPSiAhavx9M8IoXuucGhORVahsQaDTSC0gbH8NwhgSsDRHxtrEGvQg5Bs+EYrbny0oqbqYOWdTvXWfqVq+cqKqUi17zV7fiMvHdfJmeChQpZZrYjBQDL6GVXk9LDxKgzMA9qKTGwvA3VpIILsmJ3pOMwtGxWikuPqtC53f8/gRlZbX+6rbQ8d/RCxybHTLDfhbqm1a4hN4v4Vwgt9yIXGEXw3T6kzMiU3FzB19VouknytEHPrbg+vsM2gWFoSE1unQ+xVdp/ZympBaEl1+9wYaE0p8L8TQRit2U/ITeto1wZtOCGgh3Jmux0IbZ1UdCV1ShosFr2yy6LFs1m13sXItKrKHFsQy5SdTS4yh3v/+eIVof5VT16l5TJwmxRaRXNmtFNhXNsbwR8lmL0P5UfXFNMUKWaX4N9vbVrY/2R81Pj8f7E6G26zhRXdc8QqlhDqOgx6yTTycRNEPHCsIe4Q3sKDOhRzC0hs6tPJvRFoYsOT9rYWx3kZh597vvpfpfNimp14zCBMhjt9kYmB6ZVLS0Nf0yt6gIu31tMZE0fmxscTpLbRgPt/x7W3DRsV5pwhrjlGWygXVuyR40mOpPDzaJHH1t6xJANzChAvQo7/3LgVIrWAUDfX+6NlRBfZdDSqabGWUjpFaJRixZpWLKYuvPKT+0i9BQEAZjBnNLgfaEetTxpBUmUdp3tbIdgDTw8iRckrjcg/+kOUJeQNOdOnH8kx8nTfBYwNMNxiGKkn+GIfOc7hMe4mDslsOV20T5LIjtKn7DmVIMX1lJEFjfpNr4WZKKcIuzUhanUX3/X8FXt3kbuTt9r2zgwQ9gFFgrK5R+OUNolQyZR1GNCvuYRkXJTbfiFKXEAdgesfJsZymL21IA4oTMxbUNvSiYOBnKOqlVzoyKIidmNw+UE0PO5ES/i1JRMsvFxAHKxJRml4qJA6hiypKLMyz0V9ume4G3q8R8gMhHiFXyZUSLsnMxcAo5x9Df4IdK6OcqyFUiOewFWQBAE6o6IQlLykMfim2VA27JK9Xl25b5FRMqJal5C3Ghkvi/qnLoe57hxQ3YERHoLlPbqOuEorarhyAXO12M4kyBwmfscjnFvP9JU5p0I10l8z9m7EBcpXB2Ha/6hq4WD/LjpTbNDj1Tb+ElNOIKgVyDn+d8GJf4z8rehPWQkhA7Tpl4ZlHZveHieRl0p4odwYbdQI6s7QiEbNbSK5Ymhzbyiq3J4s274gQ3xcy0VpoDVSIJkhwlSeSPYDCvVqiG66JqlUDcOMQvbX2+AB+nqB25FZDVk/oUKu3N/xHix7MaKhGenQtAn1/ZofN5ejJ5HNaKcYzsM02lDazohG0kkJQwk8gr6KiwnmlqZajyoAzkn48H7IyKPOB9WArQ/cJ1xemHQmjjaF2REvfAso84fcY4dEqzzMXMAOTwnHDN+vIl/a4bDBxKT4cBINw4RWP9qC8Yh57gJ738YHI89YGZ19KhYGHJJ4Bp0etUXF/dppT4VKWRkaaxFxDUe+2KUBoDEHz+kp0gmYubhukEHHmJh4M06RxBmUBlXGvgxVirfoqE0hYPamMtqJ/LZaiQWiNMtoST4afWd5iLjqUol2+tY1fx82XuWaCiSoMUNYNmYfC92jCbHfS/rKZNZsolEZ9SQ0O2Ex6/LRhLLgRuLXhwyr0Q79nT7EV5MGlTDpyou0okOx8vpmnlU9qrLDmf366zY040Q8DtPGvlKXAjn5oaJ3v5lztSKFrleloYfZM+yJGEcx0WJOdI8CLhULnu5adnMyyakePRVwkz/apkEZXl5LjsXUC+03yfZdWyg2AHNPTbMCSRiN3RqgWHQriLfRCTRuCyAnBpRaAPr1J8q4S+0HXKEhS7zoivBq0jlq0giNR4EEcgx1JyenfKqXQcmIAaui90y1Cc72thVruYwzBDMP95FM7ecM9Uf4xQ2h2At514J0t2fQZfD6ep1Ws/8MLrpkeuYIZBS6VA0uX8JTD4tStpc49ajGXsikXh67ri5O//zDOAI+9yESczFgNUCiqlWv7FHTeOT5a/0Gkf/UhJFLPz/FXVh5bkIGIpdZcJWmp1EFtC/1apRM34WvikxkHhgpNPW9LVX4F5L5rOIgkrJWqia5Zoh2INxU1o32o1Wdao1iw4QxujBQoXV4kelrSExvi7dhL3QrBbtcXwTSnE8+WLhrlZzRD9JeZW31cjQgPK9CulX92WGN6dN8cuSAmLTPm0vlC+wcrmt8YoIQD3gk4a6xm6MBrp+gNxtFIqLsQt0gdWcDHwUFbqXWESWR60kCu5onPLLpCbemydWSDPbnOC5BR0dZ23JM3FbhROp08kn+yWDrq6EvBzhpSDujEiF86Vjx9yreAVS06QVO7UMsSJNC3nRZCEf/PJdfVWgw40p6H7AVIC4oAR5QTv5EPvRXnOwkVMLQNlqsbQ0suPcJ8pSE7ccIqhLVa5v9eN/OWd1KelmJoT0Nl0lDVL/rFeGvgNAyLcIKnM4soAyy6Dzk5Dnl8p1yRmm4+ymm1SKBV0OXtsLVry0tdMN+mmWXogumQNrVBGongv64ugBz7WF4HqkBfE9DdPSmWXC3jFlTm14QaGoIpaCE6AawQQXTUPPSbipQtlGWVmmix5JyOgR/4+jRIo2Dxdh06vFMubA7R0kY9DwzS++UYSmQi8owCzqZfAsTxBVKTFtz5UpPtuFRg+O1ZmjpstHcvVSquu2I5m8qZlU8vBRkXvlBVdVllJInclmnqXaer7dHK+QlXvMlVl0KKuvtfM5hXeaLtdrat0vfgTlfXuU5RV7F7uoat399BVvjheKT0CstHwFTKPQB29NhxiNhpg1jJC43TrxtDX6VGDB46kKX5xOC2eikGAhw+O5/moyAodqjeqnl8Vb4BdzYX+Zo8sACqPPJVjzCDeHgbLFWtccx+yWpnCdSe6n80GK0P+Ypb+SJF8omqdTEoOYhXQ1O8C6G8I2chWry8ImepsNe0fnWkC5RYDWx6ZJs47nGIw3FqTpeRcITa/Xv4pGTtgOGpINP8WdcFpr+G3F/J8Rlv57s1KT1+USREQ97Ur7PKpAYwaCPcT8D3dcMqe3U8nIkqEhupB6CJHh0Z7Xe0aB0ZbtT+p1P0SdvE2m8J1NjK70uuOfIJRd9SvtB77RmNtRXbWVeSwrCJ+cK+KNNZXRLOYKZLYfA5cfHjY0ktxUnmrRnJpw3Qp7jsx7ochXGNXseFvlT6vqkhIgpbiuCzYLBT2d3W553OWppvL5vrSXFc7d+LYvyJ77AMud9KamG4XxSoFagMYhQtrWaDq54A85Z/t+XzX1cq3C4kXjWo2v7LgBQdKUzRwWbhYBGWJQ/l6puIVeptfole4Rk++BK94bV4hn7GQqLvFyi/GK96LRD8owb7bwAd/vGsCrAS3ROwJH1oCy+GjrnIyNjtY9IBdpKVHeCzXmVfKoPLTOyvBZjTiyHReMZu9DpmVwgoeqx/g7R0NeeZcWFfQHzXRA6NzPJ7Sw8CVK2bSpbCrTxtpJFZ6jkKFxFplM/OK3exUNorvrgtfl9V5kWCBVEnzpXJkY2vV2bDskJfUCFbahwhSZhwiTMkG5C1l0fkXHBYK15hjzrNpYS6+ocWX1TolKzpHPEmtuy3uXS6AaQ+/GBiUxyYAU81sOdOi+7xgJEorm940cEdm70vpr9y8XYBOgyuNziqe81msZa2Ck7a/uyQQL9YpFi3Zu6UDlKdL6oBwp9oDGwH+PCbB6vuHWQWP9fzhJpEO/LJVlG1ZynRZvMOfn/G95R8LUHyivWIS8+M01+nv6RIR+m64dVejXlhy4VMX4HUYJq9Cj1RrzYswTmC2OQ7iZhqASm9khMfh/i54zv48OdzfTSJCYheGgEa0CBoXJCKH+7jB3aA7rA62x+HUw695nAdAeftwn8rpcJ8uLBnoORxsuxfE/QB12NbinCfhZDJF1F2KJJOnt32cO6MRzIK3D7+ZJkMxuxKECc2sHF46yyYVggHMA4REBufoQQ6YgtA73quWWa+VEDVO//2v/wct4d//+r+WeWYEYfCRROGeZa4pI81O/3ChiBWj0GQJKR7xtuU8PBYOduedu2hHYAsr8yETN3vyvFSH55n1KcjKxzS2D1+nAzrTf7PZTPnW6p4aBtcsWPvUd6nl74ZuQpJGDDjObPsw+0IWaz3UBNkbWuAwi4CM6VcBVlrokIE1gYc3YRikX9HgJJI5jzFg5yp9O6RaSchsjt4J0iFRBAIDPxDPl+QeK/sIwg9vfn7VpPPlKhJs8mPgKj1W90pN6hqAGqJknxEx5GbV1H0Cg7Gbfz8Du4Y1bS6tS6pGqUrbK/RyGYdgP7fbuAVye8/YfnMRXqOy8VstE3DW+bcwYeoSQD9GPINdahc3je/RW92lIUW6f4t9AoFGF7frxrbw4QOke/TXI/h5gb+e4q8n+OvoLf76jr5f4+8J/voBfx3Tx4ODjBT99gJSos0Sk5WvHLBSnj3J/zFs4XoFALnddkz4Y8EsZTs/VQspp7fbeKAAqQAE5NJWSV8ByjABhybT4wSQDPOgbWgaSMu8O4OXG4kapijHkdJk+YoiylNat3QShQXBMAtJOJXCF2eZv/BI6DafmENKdnF0DpRdIJ0npZ9PpoygAKRN5ZDaY1Iti80yKaW88u/Rbac4VFSoqW2MHawEO5VB6D0XVNEgqlwB+RemqJDTlEzo60o5W8/IfwUGtG8ZjmXend2JTY43L3iEzrTQgRd6e2VgWtW1/9N4Svec7BnHJ2+NdOj4ZpIM9aUwBrBLp64KNHU2I93D8YdsG76HI2h0nvUALaftma4zsMa9dnvgEcfrjpy20/JMZ9Afd/qcvXUDazaCVk/xs454+MNowb82/OvAvy78A/sx+vBvUAf9gKnTz5cc0AZb09Xo3lUZ90ddb9wfmyOvbffbDmmbtm17pk3scbtnaavyn+yC8O/XoBNCT9uUeiL//tf/MZsd2/jv/2827T51TMymaXeoc2I2B4PcPek+uieP7sl/qnvylA192TV/RpU5IEloOHQ7Yu2+jskv8/nNyUfnGbnsuy9+C3fenvz96uQy+vDsl50JefNhx3r1HTmJb568fNp/99PT/q77pH114v/jaPx05+h7t9+/efXh8qT94uWzp7uTl7N3u/aldQX0fvzbL7v9618+7F69eHt80j/+rn25nOxMjgaLZ0+sE2v0+snTnZcfP4S7i2n/1Un3+Mn3T/uzuf3X3ecfvCfw3v71aTe0Z+HO0Y+vBiejH26+v9x9uzyJd6Yv4p9Odp78aH7s/Phy4u72f3v+6qTt//D2svfuaNfc6Xz3/cVJ8uT65Gm7/dp+u/OP2buLk/785s3lrvvr9V93w2fHfz0Jb1rRL92j3qu3/VeLJDjZfXLy+nLnpfv6aLf/80/tk8Wxufhov/zpxt19FxxdnbSOXr65HLxbvHu38/Qn6x8nO8e/OZedlx+74U7RQeOd1kYu2qf/07p4NnoIyJgFD626wZy+7u/t9AnU7CI1G1jRE2tRByTDtYq41p5h63Htu4e5mvm9mbmz2YCRot9pdVp9u9fudTum1cl8T21W7omyxQDJEaUDkN23u7bV6re63e5gYCuOqZao6KaeDrrscGmrDfU5tVpt9tq26GsXPI5WF//R1z489qCMNnsd9NEPwS8+9/DdNvG9hc6J1aEJFib08LPiHZZgQ4JtgsFYA0rCbuMxVyzBbjEI08aDrvQXK6SL6AMcn02TJlh4HLZn0lRM6HWQZSypTzFa+Gi3snckNmhn4FhiFyvdtSk0UGsjRcs6+xP47HVjXTFWoRhbLsX+DIXYhUJaciGtx9nHZi67NRq7dqdjQ/Putr1BdzTwxm3L7o67HcdxoXfNXHakd8i+NAnu3wwD4OcSt3vg71abkmu+d+HE1UPmoDZ1TjvF2aP+PfHwvq4a9b2muAr+pUN/wH+B2tHS/7h5gyZLKwOdNgF86sxjzKTTD41BKDS+1E5B/oOnKitlUm7Mv/PMhk+XT0+hVmj/PcPAr3kNev1ObzCgz62WbQ9gdDjLd2IAdLvb73b77S5CdLqDfr/fo9BWt9MxW2brDMBTeKTeskxr0O5SmG5/0O6bnTY+99pdKLbTk6nDiGda7U6nng+/FDMdrVXq3bZpti0G0+na3YHZp7h9GM9hZDdk6oNWu9PqdVoUotW2YIZOMXt2pwtjmqFS75jIg8U5ACfAtijvZgvE0u8pkhl0+qbV57zAgzXg1Ac24MGQrFLv2DCGM1l3LavVsdlzf9Cyep2Ownuv0x4MzAEt37asbtvqDeizDbUa9ApyN7tQfstsMepAs8d00B5YvVbX7MjUO21aV0odu8C+2aWYIHUY3a2CZPpmCzTSYtyYnS5Ip8eoQ0UHA4V6uw02MBhQCKsPUm8zmdp2a2C224OCVlsmmFmPS8/sg6LaTKo90wZnTaYOLPbaLbPDbHZgg5T6TGMo1V5HpT4wO4OWzfTUAzPs9OwOs5+23bK7itzb/a5ldtvMBgZg5a0OlTs8mN2BXdBqFyzG7DNZ9/oDlBKra7dlYaky9X4HmoM54NIA8j3WslqtVtdC6LMsVMU7sVpJf/FpQ2Gn0xn3XLCr/hga66A1Gplja9BpE3fc8nrtx+jVY/TqMXr1GL16jF49Rq8eo1eP0avH6NVj9OoxevXHRq9g4jLo260+6cHUsNeHPsBxO/0ezMBglm2POn+8y/4YOHoMHD0Gjh4DR4+Bo/+wwNEnT/DZXccvPPqhlXxOaludUbfVG7edvtMejL2R63lu3xy1XIu4brtHZxwb7w2TgNftvhKB1677iMBrI2Mi8Nox+WzIIxyhG4Vhoo1xpGEI1AIH4uBK4GG7mSnqHLO3C6jH05AeX0rfmy4m0GLo15iHeU5EwNtwya9+clGV0DOi48iZ8LvRlDMhT/nrcw6BdUihpXMFKWHM11yrmJmNcq0iPWKdpPeeZZThlZ9GeXLzAoin2PjB0yHFicNF5JKn9ELhEhl+iU7VtrFjKOj8um5RLBm15tiP4rRsWjNAyHPzmM+drAmtkNeFeP4H1YGcIQ==</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_68b34aeff02e43029cf4dd2722226fb5\")) .filter((elt) => !elt.dataset['step1']) )[0]; root.dataset['step1'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "class Count(nnx.Variable): pass\n",
    "\n",
    "class Weights(nnx.Module):\n",
    "  def __init__(self, kernel: jax.Array, bias: jax.Array, count: jax.Array):\n",
    "    self.kernel, self.bias = nnx.Param(kernel), nnx.Param(bias)\n",
    "    self.count = Count(count)\n",
    "\n",
    "weights = Weights(\n",
    "  kernel=random.uniform(random.key(0), (10, 2, 3)),\n",
    "  bias=jnp.zeros((10, 3)),\n",
    "  count=jnp.arange(10),\n",
    ")\n",
    "x = jax.random.normal(random.key(1), (10, 2))\n",
    "\n",
    "def crazy_vector_dot(weights: Weights, x: jax.Array):\n",
    "  assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\n",
    "  assert x.ndim == 1, 'Batch dimensions not allowed'\n",
    "  weights.count += 1\n",
    "  y = x @ weights.kernel + weights.bias\n",
    "  weights.some_property = ['a', 2, False] # add attribute\n",
    "  del weights.bias # delete attribute\n",
    "  weights.new_param = weights.kernel # share reference\n",
    "  return y\n",
    "\n",
    "y = nnx.vmap(crazy_vector_dot, in_axes=0, out_axes=1)(weights, x)\n",
    "\n",
    "nnx.display(weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "743bcc34",
   "metadata": {},
   "source": [
    "> With great power comes great responsibility.\n",
    "> <br> \\- Uncle Ben\n",
    "\n",
    "While this feature is very powerful, it must be used with care because it can clash with JAX's underlying assumptions for certain transforms. For example, `jit` expects the structure of the inputs to be stable in order to cache the compiled function, so changing the graph structure inside an `nnx.jit`-ed function causes continuous recompilations and performance degradation. On the other hand, `scan` only allows a fixed `carry` structure, so adding/removing sub-states declared as carry will cause an error."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d11d191",
   "metadata": {},
   "source": [
    "## Transforming sub-states (lift types)\n",
    "\n",
    "Certain JAX transforms allow the use of pytree prefixes to specify how different parts of the inputs/outputs should be transformed. Flax NNX supports pytree prefixes for pytree structures but currently it doesn't have the notion of a prefix for graph objects. Instead, Flax NNX introduces the concept of “lift types” which allow specifying how different sub-states of an object should be transformed. Different transforms support different lift types, here is the list of currently supported FLax NNX lift types for each JAX transformation:\n",
    "\n",
    "| Lift type        | JAX transforms                          |\n",
    "|------------------|-----------------------------------------|\n",
    "| `StateAxes`      | `vmap`, `pmap`, `scan`                  |\n",
    "| `StateSharding`  | `jit`, `shard_map`                      |\n",
    "| `DiffState`      | `grad`, `value_and_grad`, `custom_vjp`  |\n",
    "\n",
    "To specify how to vectorize different sub-states of an object in `nnx.vmap`, the Flax team created a `nnx.StateAxes`. `StateAxes` maps a set of sub-states via Flax NNX [Filters](https://flax.readthedocs.io/en/latest/guides/filters_guide.html) to their corresponding axes, and you can pass the `nnx.StateAxes` to `in_axes` and `out_axes` as if it/they were a pytree prefix.\n",
    "\n",
    "Let's use the previous `stateful_vector_dot` example and vectorize only the `nnx.Param` variables and broadcast the `count` variable so we only keep a single count for all the batch elements.\n",
    "To do this we will define a `nnx.StateAxes` with a filter that matches the `nnx.Param` variables and maps them to axis `0`, and all the `Count` variables to `None`, and pass this `nnx.StateAxes` to `in_axes` for the `Weights` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d10aee8a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Count( # 1 (4 B)\n",
       "  value=Array(1, dtype=int32, weak_type=True)\n",
       ")"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Weights(nnx.Module):\n",
    "  def __init__(self, kernel: jax.Array, bias: jax.Array, count: jax.Array):\n",
    "    self.kernel, self.bias = nnx.Param(kernel), nnx.Param(bias)\n",
    "    self.count = Count(count)\n",
    "\n",
    "weights = Weights(\n",
    "  kernel=random.uniform(random.key(0), (10, 2, 3)),\n",
    "  bias=jnp.zeros((10, 3)),\n",
    "  count=jnp.array(0),\n",
    ")\n",
    "x = jax.random.normal(random.key(1), (10, 2))\n",
    "\n",
    "\n",
    "def stateful_vector_dot(weights: Weights, x: jax.Array):\n",
    "  assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\n",
    "  assert x.ndim == 1, 'Batch dimensions not allowed'\n",
    "  weights.count += 1\n",
    "  return x @ weights.kernel + weights.bias\n",
    "\n",
    "state_axes = nnx.StateAxes({nnx.Param: 0, Count: None}) # broadcast Count\n",
    "y = nnx.vmap(stateful_vector_dot, in_axes=(state_axes, 0), out_axes=1)(weights, x)\n",
    "\n",
    "weights.count"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cfd87e1",
   "metadata": {},
   "source": [
    "Here, `count` is now a scalar since it's not being vectorized. Also, note that `nnx.StateAxes` can only be used directly on Flax NNX objects, and it cannot be used as a prefix for a pytree of objects."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c8bb104",
   "metadata": {},
   "source": [
    "### Random state\n",
    "\n",
    "In Flax NNX, a random state is just a regular state. This means that it is stored inside `nnx.Module`s that need it, and it is treated as any other type of state. This is a simplification over Flax Linen, where a random state was handled by a separate mechanism. In practice `nnx.Module`s simply need to keep a reference to a `Rngs` object that is passed to them during initialization, and use it to generate a unique key for each random operation. For the purposes of this guide, this means that random state can be transformed like any other type of state but we also need to be aware of how the state is laid out so we can transform it correctly.\n",
    "\n",
    "Suppose you want to change things up a bit and apply the same weights to all elements in the batch. But you also want to add different random noise to each element.\n",
    "\n",
    "To do this, you will add an `Rngs` attribute to `Weights`, created from a `seed` key argument passed during construction. This seed key must be `split` beforehand, so that you can vectorize it successfully. For pedagogical reasons, you will assign the seed key to a `noise` “stream” and sample from it. To vectorize the PRNG state, you must configure `nnx.StateAxes` to map all `RngState`s (a base class for all variables in `Rngs`) to axis `0`, and `nnx.Param` and `Count` to `None`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "33c284b6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "False\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<script> (()=>{ if (customElements.get('treescope-container') === undefined) { class TreescopeContainer extends HTMLElement { constructor() { super(); this.attachShadow({mode: \"open\"}); this.defns = {}; this.state = {}; } } customElements.define(\"treescope-container\", TreescopeContainer); } if (customElements.get('treescope-run-here') === undefined) { class RunHere extends HTMLElement { constructor() { super() } connectedCallback() { const run = child => { const fn = new Function(child.textContent); child.textContent = \"\"; fn.call(this); this.remove(); }; const child = this.querySelector(\"script\"); if (child) { run(child); } else { new MutationObserver(()=>{ run(this.querySelector(\"script\")); }).observe(this, {childList: true}); } } } customElements.define(\"treescope-run-here\", RunHere); } })(); </script> <treescope-container class=\"treescope_out_fb24928ebcf34fd283295e04e5f0c2f5\" style=\"display:block\"></treescope-container> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_fb24928ebcf34fd283295e04e5f0c2f5\")) .filter((elt) => !elt.dataset.setup) )[0]; root.dataset.setup = 1; const msg = document.createElement(\"span\"); msg.style = \"color: #cccccc; font-family: monospace;\"; msg.textContent = \"(Loading...)\"; root.state.loadingMsg = msg; root.shadowRoot.appendChild(msg); root.state.chain = new Promise((resolve, reject) => { const observer = new IntersectionObserver((entries) => { for (const entry of entries) { if (entry.isIntersecting) { resolve(); observer.disconnect(); return; } } }, {rootMargin: \"1000px\"}); window.setTimeout(() => { observer.observe(root); }, 0); }); root.state.deferring = false; const _insertNode = (node) => { for (let oldScript of node.querySelectorAll(\"script\")) { let newScript = document.createElement(\"script\"); newScript.type = oldScript.type; newScript.textContent = oldScript.textContent; oldScript.parentNode.replaceChild(newScript, oldScript); } if (root.state.loadingMsg) { root.state.loadingMsg.remove(); root.state.loadingMsg = null; } root.shadowRoot.appendChild(node); }; root.defns.insertContent = ((contentNode, compressed) => { if (compressed) { root.state.deferring = true; } if (root.state.deferring) { root.state.chain = (async () => { await root.state.chain; if (compressed) { const encoded = contentNode.textContent; const blob = new Blob([ Uint8Array.from(atob(encoded), (m) => m.codePointAt(0)) ]); const reader = blob.stream().pipeThrough( new DecompressionStream(\"deflate\") ).pipeThrough( new TextDecoderStream(\"utf-8\") ).getReader(); const parts = []; while (true) { const step = await reader.read(); if (step.done) { break; } parts.push(step.value); } const tpl = document.createElement('template'); tpl.innerHTML = parts.join(\"\"); _insertNode(tpl.content); } else { _insertNode(contentNode.content); } })(); } else { _insertNode(contentNode.content); } }); </script></treescope-run-here><div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtXAtT2zoW/iu66eySLMTYsfOkMOvQJNAWWggtbXd3WNmWbTWObWwnIezw3/dIdh5OnAD38ugDmCkgHUnnfb4jUV6H0dghe0IUEBLqnk8uAs+L0P+Q74U0op7bQAFxcESHZAeZnhsVTdynzriB+p7rhT7WYXxk04gU+Q8N5Acw4tAwKvKti9HYh1HXc2FYw3rPCryBaxR1z/GCRrx0ByU/aQ4QwH7UiOwGMmkEZG5E3GgH9albTMYlUfwb7OVdFUN6TV0L1nmBQYIiDO0gHxsGDBYdYkYNVNJtxo1Lijahlg0jklBm57kRpiDcdP/km+KQhlSjDo1ARDyIvCltkbpRQN2Q6uxYEs8mct283o71+Hqqx2IwcOHMAMZCPaB+hJgidjew7ztUx0y1254eEaamgOD+xl4+X9jdA83DeWGEDGK6IdpFkU1DwSLRKZjl2DNIviDYXhgJfB5EIxG68InLRFZ1titb9K//ZM0cYNdwCEy7A8fZiU8QgM2u57kwmh95Qa+A5nnwzmGITaWGI6qzQZ8Ephf0sasTwfVG+QJ3BDggvzSDivGi10guFWAfaqL8AteCQ1wrstHuLhIZyVrWAxINAhf0jogTkhlj9sBlnC1uHdrUjBh/nIB9cwOfK07Ig/u5hjcSAnI5IGGkurTPzdUOcJ/kY50U2B47Swf5g9CO1biTIePkiN1YjDVS3p0HxkVsyMizLCcO3wseYuCtPtuLjRAn2kJkCA6eWJJxx38WemTMlJ4LcoyhhFjQHRyG7yGKk33zuemeF31ww9zk8JsC6BPcn/v43uvtrAAw6BDxDXdz6TyTQxHWQFJytZsTcxC6QbRM4rnAIijDhal1wZCtgTxbM5E9B8EY5zvNiyKvzxJDw/WivGB6joE1WO3Ctg0bh/k9B2vE2UvPXMRn8DUN3SZ6jxiFAvoHU90k8USe30CiIJVJfzn1sLHrIhcZft7JTLM3Ak+HF1jTAjLk3s2z46tKrYRFcUage/0+iDVHgfkHM80CCY5Ztr0hCQoZ9Al5yFKcNb+hLJWl8pSAGBdMohlBkrGXE/sQB/liUXM8vRcPFXYmaZorSPKvUOg51FhHGVvpNuIbZn8SMOZ8YIfYYDESzAtRZZ9J/eIO0EA0wpCH2eKUgVdUPtBnlh8AtUFDOHQ8qXCLhGgPcT9qNDQCOXFOc690/pHtAnHxKkqseiVVD+w+PYu63K+4Gtacya29fLKBg15IsAVx5i6vfqB4mPLAlmYvmtCnOORVuIE2/l0qa/rGc7KXXrSSycoTMMnsyA4eBCEzoO8BFiFBxrk0fLhjeSjwg4o8e4arfPxhTp2JF5GraPkUgYYXJg3C6MJz4yS0HFrrQkkolVk0ZZoK/WX2Y4svssik6uPAAvQYs8ED+uYvngap2B9rA0iNbmYCmk1nOW0O5RaoQJGA3bOJ/00kxcgtQP+NIwxeQbGDuuO+5jkh+jCImLwG2o9Xwld/DIFRHBGtBzA+zrx9qMw2B+zYjWA5xSExpuD/FRHZ586ym8erOegWhToroWkp4/jIkCI73c1WCiMcXuhQB0Cx0/XYjFLVY5Kn1525sCZ95LzqkwJm4AgXsQuG5bCuMD/MDmFYNcDuxJv5tkgKEQGNAXooeoPofqJMOQDDUGL8keaEH4n+oH3fCyLsLu2tBV6PuBdsZJaMbtfu3LI5fU7MfCPYun3Bu8YLFh1zECoOFlGII3aOLmBAao4wiIHVlJIBKxDVuNChEzEC4ibCp3s94DJN+GAYcBqMSWlOBb+OHT0PDSl0QwBoOIAWwgiz9VN+H42TBFHFnBiAq8BNGKyaM8flADsu9BcX0MKb9Ao2SQVejQceNFOYIa0RDlywwMWkVEysa5pYl+QMQh+6kSXD8XyYKCkZKibGnEHKBkfROChaATYomC2PJLlsEGsLeRAkFkEisFfR7a04aKCNYCmID6FEzUu8LOXqhykEaCnlF6fOuWjrVb65SDeF3rENIbrTJE/uvaKg6HbhLrw+jTvHrQHnKrNBSHVmXA6wiWODtfISiDO/JNk3BCpSYB12ooPZ0uLd107OnbQxfOV0y8loMRl+Qm3OGjGuwSKLq0GYCCiRPki3WrAMDHVvBHObi/8v1SvG9mXZ6g694nriGwF2fJS2f2bkv9L5j2womKzxz0J2DGiD8jJp4pZwjgp4XbVRvAOoxME+ILDbu9j7G3j1CTNGYyJyBVDFWEnzEHxkHZFSxey+LEsV6QspIX0RhtbssCiqm7qayya8y2Er93nA22h2QYj+UIMAjwUz8Pp5w9MH7CpJYKAgFIbYGRBwtIIQen2S51CBXS2yr0LcKLBrxTu2CrkNKJyF6UVuaBMSsdteMkL73W6XSdNlY+zulk8KAeF3Pd2xq+f/+8+kPdHJBLTcv1WZvxxy2e21k4yNkhhW2NVbGOgNNAicPMPNDTa/PfJMs7SjATKvKFuGWO8cWWpT5R+HJ6rq8e+apyP496Ctqi113Uezr6pWz3tnHLaa+6Ovqnr2df+tenTY3Ffb1tXhwXs7CptHlFhy+82X0vvDytdh1x/Qj0flM+ntl8PTz0fD86Pr6OO43d7fPLd6Z7T5RrTpm5PB25bR+S4eaNvm8NDwL99V7MtzSk8GR27HPjA/ReqnSvM4UNT2odtrVfRPg4G7eVq+1MPeaGi2ne3LK6vl1Szt7ahTkw7UbVc9Lb8PgrfS6aZ1LZ4aovrWlKzj6v6o871kid54cFqt9ltSZXTwpf7Bsnxy1hsr5FC7Luta8KETYdU6OTwevcHhODwZHB5+OW+1R+rHE//wq/Fpe3vTqp5Vv8iRaL77eKkOy7Dne/W4qh6N1L51fdrdHHzrktaXq5JZ0a+PldODcXnQVN9dN7/7bV+mByf7LfHb4KPSrbpm833roH3UV+lmbdgq2a5kVze1z6Mv30cHwfBN59O++91staxo84P+zXGq5fr+21GzZteVo6NOV+58U63+Yfl786QenXXIQb3VbB525DeWcrr9VR9ragds+vndtnrSwSo52nfUg+vWB+tbZFWaH60PHw7fNHv0pEzazS/7zbZORd8OPN8F3/C/td5I11Kva+6bkT1+5x4YuB0emOJxv9M6rjQN9fLzZx9HYfdb3zAwrZfM67ryiX6/rPj9oPLB+7rfpUGnP3zbkbvnXbndKunNE/Ns88Dx/I7SDkdlbF1WavQb6R47/rnbPDgkxlFABueXnf2+dN4Oet3uVblUOT8PRypwVED85SbKb3C33mAl87/wzzT6seH50DvMQpK/NwmCsIZiK47Z/8Be62/wbf4AwhvGuJeFvcE9XB3l45Yy/TwFIXjmsfAFsqTlZGMhpAe2BeuhWeOJR5hGyMVDauHICwTY2dc8HBjCKKAROSNXUX62F0MU8V6zNxAo8fncXIPNXj/glDPaJ9CJ5yfPY0vrAtKHbnlp6c0WKomiyJEUJF9ArXl+VZR97lwXnZsxxy7JJhmMPRjl0CvUxtSBxBZ5iBH/wTMboE0X+jrIxhR0RrDBLgE253WXvOTc8obDLhQQT4+7uRSoaiCv5+g2QOpqrczQjygjqVoVSvWqXJFLNalURtsAgEDcLDzJLrZzyfbJG1H6MnuxTwPiGAC8pq4/SApZjpd8zbvKZW6SoAOYjJEByMgXp89NvXAkMBGl7ygWOE2X/tze352I6REo1tOlJhfabphNdoi/nPPyE+ano4n0S/tPUE4uQ6rJO1BuD9zkIw4wANU6yssV1CxsgZ8M2OWQhPIK//nUtboR1M0GuCm0OfAPGz3zIuw0kMyGKmwozef8oRsT8LmRHmZDF87QSQ/nlhrjdDeQW0O9RKpRHO6udqe/4j3ZXvM4PmA6+Epw3SvAWYBegBOHakJa4dyM+Sx/2JjdwW0gz91nmWd3456plz+cFjbQ9IJwNycw5eZmWWA6BTiLz7Eb38W7TZjjaQbKmA3fJ4LvrXaee3qzzDz0Ib1x6Q4xt24+t8eR8O4dbE2N3dz0pdKolOVKFWOpXlIUpVLBWrWq1Os1oy5WoDiUFg7NeuFMfC1b0u/gPxy9oyyNTFtITjghMR0PR3IpL28V/m7NvPgWS6W6kUxZNUOvSZVKXZFlSalVpZouK6ZSNoiERaNWVdbl/4cJ2Kyk+eOo9PYY2LpvnKzLNo+aHuLGcF2SiCkeJlUspgPIB59pCMmUXvMdkU0NiFbAPmjWMLNO+k859S+hvrt/WeHd6wywqvjN/0JLbg+6hTsYczleCusL7xL91jL99EoKnOWeMfUTVtXVTr4eROkMEP4CKCqtAA5znxYwcUWusi2ffA7IlAD9HwIxpYBCTSdVs1yWZFnDilkuMdQgE9GsiYCMNAk/ByiiLqvfWdX7Aev14zrg+qIyR/KoVeWlnPz8OefPFpQeCVzivPTlj2TyWL2rbB7PPkehqaB8Sfn5enNdKclirV4pVyqiAiWpJmKxXpLksm4SycDV5+zNS1tIftjuXFFEQzartRopVxVRr2taWTfkUgkbmqTpov4bdOe3KPVX6c/jRLAeDczTvPToz6jCF2D1owGrx62yfxZaBa71WA8e65/LHhlggVzL2OoUhH1aaMXUu8rkbO45YFXW+9yzQSzXoyH5nRywy39D6um9UOCaXueLMcVv75G/yu3lrZ74DBeaMzdbf8+wSPcs15zgior487Wf1ZpWK9dEo1SSFKUmm5okyjVFq4qkJEqGLj5H+zmIr0El8YGfhqtarS5qhqTJ1apSVSp1IlZrZhXkxnKpjH/lp+E7qfQxAPLzNJ+LGWF9D5VN/dKQ/mBqvTUNLGbECFu7GcaI/+4AhNirq1J1hzMYf/szvbUsKReEvbMhgPbl9eUXuCR4eoB09+DrkfHvAIvfkfGzgWJQ8R0sDlTPBYhr4g/58q8T9ouQuiZquqZoBAMolLAi4Yqk1MSqRJ4D84KV2I5mMGbZaAVQ+/nSEoh1Z5AwpX2K9PSbYAOm07shg4TyBRf8WrjgCWrES9At6PxuAfcSbLcE209m9se8LH8x9w9k7sd5ofsZTbz6fzTO/uQlWviDSrn7O8bdW5IJpUGHe/8H8bcHcA==</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_fb24928ebcf34fd283295e04e5f0c2f5\")) .filter((elt) => !elt.dataset['step0']) )[0]; root.dataset['step0'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtfet64zaS6H8/BeJcJLUlNam7LNvz9SXZ9G4ufbqTmcnx6nNTJGSzWyI1JGXL8ej/znvseYDzCudR5klOFQCSAAhSstvZ+WbSTseWgKpCoVAAqgq3E8+/JnFyu6Cnh54frxbO7TEJwoAeEt87PZyH0YVH5zSKqHcx89yRPRiMe92u3RsN7ZHb7c17fY/ajuWNhr3Ds5N45QTwG+mdtZ0ocm6v/V8v3DBIHD+gEbkjN1d+QlsA51IsKFo6iwnZEhNw2w/mIaDMIaU1d5b+AnhbhkHIsCfEDRdhdEw+d9jPhCyd6NIPWrMwScLlMbHanT5dTtQSVxGtLs4PVuvkPLldgUAiJ7ikh1Ng4ZpGie86i5az8C8D4ML3vAVQmvuLhAIPl0AthnxatxskhKL85LZutfuNexd2fBVeM0EVSd+PXrBezmgEBIMwqR/PQ3cdN4DsLIw8GrUix/PX8THprjYfR5J/Zkwj+bRNhuxnIoo7JvZqQ+Jw4Xt5VkWp7RggaRTr+lLVeoyFxF8BjqLIE7IKYz/xQ2g2ZwY8rBNImznuh8soXAdeS7DMCjIxPFsALFBxPM8PLrleuVdI1g+ghVr0mgZJnBZ243vJ1TG0XtJC5iBrQpCz+SK8OSbXfuzPUHGK1fq15Qce3UDJlmVV13IWbvasZbhpxVeOh0Vb7D+sFqtQUyR0IEFU3VyhjK9xBVvuwnc/eE7i3KfFFqGDEr1Y0jh2LqmkPWmP3p485WPJSRJRGrvhiraiddC6ohGkxW7krxLCdLPmrFbAg4MSeBq6CU1aMeA4y9rZAf5AsXFCUi7IKanXG+T0jNwdEAL/5uvARVTi0ZhGPvTyX+nPII1RHTsNABAS0WQdBYSlPkM67XkULutOEs4AqEnqS0Zw2XZDj75GUT5L6lajMQHs7UF5Md+AGJJuJy+Iszq7TWgMfD6ovJTIHGkjlYDeEFEQo1Vn5Nuz9RzGdoEiKshxdnH9Kvgf4ZkVcy+OOcsLmpAXqExLZ/Xm356/BM2c6LW5pMkLUEY/WIfrmAHXr53Fmja5GgImoqU1RIozJ6YXrDc0STifxzThfPhzwlHJySmxUgwiwUN1rIlI5Zh5ypbQRUwlImenxC4hInPWXtDgMrkiLdIpkLbbKvGUGBexm8QZRV7kE1I3k7YbExMf3zvJVRvkDjLLiDUKXOTlfElswY/U0pFWofO8iOm5NUWmbGCBk2uQI0GelCGRI2ILRLl1eGGXVYXZDy3MNhc2qyqs89DCOnphQv/Poya5bJLZ1NxpbwMwndxnkRv7wdUbCtTrorwP9JaN+X8Uar/wV987oNuRc/O9H/C/+F2Q+DdnlaplRj1OHJjP3qJ94oki6oCROOtcg1GzP/Pjb/wAJoY6y/rrX7kKwVRV3zTIU0QgJ8SmrV6Ol1Vwk2qWps0ZAKMVg2WGxJ4wYk8yGPxhAIvwsl4s9Uhg/yVKoFHEt1V4U99wgCbpNBqZbm8lLc76viJHcqqMAZjP5allaOJHjeHyL62sDs/JmjjTIbNaKxnpKMpaOYVaOpt62u6CocakpKYnGUTG5T/oD+920FrYZ8jpgdbydLOqZyqgygC6Ya7eoDsZWFq1rOmzMkBUaUlPeUdWxfJUExyfPECwMNYCrqnLZF0Pe8dEUp0US1OdjBj0GNtSVSAbXrGrryiOsqjtMO9WddangmZKS/SslIhxNvno1sqVPpVk/tWkqjohlKtHg3DpB2BjRJkO+0FdUgFTtbWhT4iAsSCNds0dVNKOguhKsylMaW2nMmxsQCH8vZpNIpdPsYVZIAnfoiH9NonA4uZzvWq7ZZNIOnUpU8y76HJW/+Iu2pIv7i7x12zbeGecbtC+j5wYzKrLh5UoQaA7EoBXcwsQ7Y7dGUD/jGCIbg/tfgc+X+Jna9jBz7N8kMrRzojdGeWyF5WpMQeoZlTpFIQ5MTVFoFpFwT9fvHYScJgCsIGgPeDfbRPcRJaUTZTQyHVUU58ZfPDnJAURRhakHR01NAstCm8AXgCe+9NUQTJy7zm590AOYDNS72VSYvIJb87fT+VUKCTZtJH/N9RN6mhdvAfe4Y/fJHZTsvlyjdwWVIsz6vmXfsKM59eRv3QibKpzBlv7fM5+ak34aM+Hw1mPfZzPh3OLso8d17E6LvvoDTrDzoh9HPcGw5nHPo7c/qA3qzUFQdodDt0Oy5m5M6/DP9rDGXXnNYBhYtL5ekshxVM5G87xP4btUHdIR4Kz2WwoeBh585EjUsej8YB9dPszy+vzj72xO+5lnM2Hs4HH2fFm3mzE2R9Tz6H9jLODjDuXLhZvwYsCloYTnqE5LeCZzP3Lgs/iwYjzY0BfAH46wjHdg0ZmbkuTCBfGj2Fk873cl+EEm9mAniqEMM4YtKSFQkFYFwYua+GHhXtV7/e/xKBBozY5MCgSFAUdcci44R/wX2NSTXNoaTQLHUvQRSM4o82/nGd6em41Sf5v2lQybJZqFzMeBWPaKPhxqtDbGKGAaQCq7GYOZ03v8q4w+xXrpUAqhVJmGz/+wfmBO48NuaMXJC4PfPdtPgM5nBQ6/T6TB/xtSJQf1IggYlsSsT2V59/SZskaxi40ZTVWWVnThmaC5645tOEraAdwZG4VQfMWxKYCuRgCC+CPN0mh7YoiNU3Tj9VclTPzx7ccF7F9/+Z5WJaxrOqWaz2o6azfRdPlEjU1Qln3qe6P1V3V3tHp9CYSE5eHjq+Yw8SQN1HtnRTuDKNgmt3z0H563+a+b4M/uMkf3OiVHWxXpi31P/temOWZU9nqzHtv1pwnxLp3c1q/v+a0quRuPbg5703W0Jxy0+VmS9rADbVtMx///raHscACCSUPf1QVMGhTyqmkUkrTbR+sL5qXtcuQJLWVA85CIqLDtUL47aMsSrVXpdnnbKydymanAAF+1sAGTK3U+2R9Pr71qU+EJXETk4yysE2TxW2aUuDm/k2wt3Iy33sWluolOrGvnSiJn9++RNDMMWdykR0ulN9wWkyBZus0SdecA4Lt74QYCIge/u2CAu2E7AtIhtF7CEYX//Y/DhP+DlBfcu9TjjH7wTX68SDQuQOtJHdWbgV/SWzymRaPzK2nDDuJ1nSHFgb00kn8a5otIZ7kK5wpzNK5BHN77SlLELrNJsVrcNU3w2njeMyG4kY7Xi38pF6raaYeR0oXK08KiiVyTEYDgl6sEBbK1PHOFcJTeYxHKa8itkPiIqIr6iTxRTjH1er1YqHM44bAn0J2Qo6OfH3OEz08ToCbJol9jwoeBJecZSkgWJAhAP7AtukI6QBsQwXOJMeG7mJldJ6Kw4ohuMZFprFF9hNVyXxdVa4INt6nVAaqFWu0D9hYWLAO8hlENE4+ixRaS/5caiFoxlDanxrVVktxFi2YSmWTnzSun2vVOzfbfrmTZ8wwGITTkuoaGliuv3Hgkc0iOdjFgT9ZGo9paRj3iuwthv3FW7KUgmYEX4p7FnivAs93aVzXI9k+T8cPMegJQ9U2IKUD+rkYEtggnK4lARKBkcCEnWoW5JzX0I6pTZklA7MOjZxFTVY2VkZ7tY6vUgTGaM0YeyqS1C32lPVs04yo5rlAdTZ+XJtO1YkvBT4lAir+4K8u2DhU05Z6JHbffXFnAN8eq8k08CDxndkbFwWf3LdcjmeiuguHtExl4c4kXK+q5nuf9uCNpy+aFRZEiuwxxHvUqFZfOvEH6pFwnTRqD2LzYhGGH9arArfp+g356ivymcD1L4MwQgeRjZYVrVPOV7E6XFXj9SxOwEZjfTdTQc7bBVurrk01dzHlVEUt8xwLHK6DD0F4EyjslVgNEp5cWNm8tI/ssQuaRI/WnbnbthHF3GfPdnedjOZDekBRr+7VbLxIjfV9W21nm1X3kJL22pY5InJxJ384q2lORbigbRpFYVSv/cx5kcf+mphHjDu7xC4AXsD70A9S30PZYPoMGvntirqFVdoLMP6c25+DxF/8kW/4rnsUY39se3KTOAwslZ00f7Ht4T7MOrc/zmIaXbPNO2IfLI1iyvDSrHodbNzIp3G2kTltL5F+bk3bvoT4BosHBbSKs4/Yiv29E33AbfenROK3/Zc1jW7fgj3rJmH0bLGo1/St27LoeeXq8rJE6grRBXYYrTBVhQCkHdFleE3rDZMqFyXU9vwYKhGg7aE3ZpPcbbO9wlCNOHkWgOOADH4TOUsqbQIvIR7yD3L7paaMeWf3bO0vvGdin/k3/uU60hrfZdGStNa7VEVl8GJf6iqLsmrek7/UeJqHPIiNgR78Ju92ZbvqnzsxHfRyICmxAPuSh4oUUJYmQ7J563uYmnTKWoaMgycQMGjgsW3oAl5KlGE3afeVQPM0GfLWAHlrhIwXMAd4BnAtQ8ZRo2g5iqvtBJGsXH8DI+FrGuE2kBxBSdYkuabfMBsb1x9eSQawItUyoMlBPk5hN2btlW+3kpsUZss5P2YgzZQMgI8p+qEHSUuKuyg0wn6wkyw/lVBKNMVMrtA3wXH1az4/rIN4vVqFUQJmkMdm/kZxuzrTuws0ltRC+TkRTSsbstAyyYULdHLFVnhMcNdRBKO1mhjTFXNiLNmL0UJJ2V7dXGXbafDhVk9qaHvMsAA+b4ry89h/xg+yepR+z3b2M/6VxK1eT7XGSZg4ixfhItbqHS7+hAelWD3taZ7BqwNSLbhxmgAK9Tbuo4sBAHByYCl0JuKYuCMXctpgPYFNxT4yEysFyzniIstNE7lu+ecnQPKIiRjKwp3+gXKMIqs5p5bhKQqXSe5NeKNJDjT3W+pfXiUF0d3uK7rb+4ju9iNEd1stOlG5/PMO0eVVl2SHiI0SVRRTjotxtbc4Av/IrFss9G5bLh9tsN5DSBqGJKli4edMTgFYH1PpDNCBzrXrgOcWy0t3dS9010voeG03ok5Cv15Q/FavcdBadoyKfW2zg4i49ztXzSPSweMR6e5DBfyKSTaDZ+1hhpeE+oLh4jop3SQKr4KqWEWF3Hqt4zEWDYOE2HHNzqb8h9iYrcSJ85g9mJ84vWETqqaRtAqSzpxv8g3wcsi5aunKuD/PiCJ4Vq1Xsc8fCizAL6XjANLm/sLucLG7X98ZXlhf5nti2AmQZ7PYVGKWqW5MyBCdTQUiyywwbG6oUy7f4lpGoRnEhna25RgFglvucfO2Lhzz+r0wINkG+O9TkhXHm/IfE99NLLWZi7CZC6VJrLbdb0z2rY7CEyYe4QG+p/m5vJ3BcLFehoeEDKrjB4YmvI+e8cNHS/mgzq5KWe3+Hi1S0cItrA5rYeSUf9sRhDkwnO4A4/2VfoJSCuuyARh8Sm0EVs5bZCSOTmUPQR6In1QP1BOdQ+FC5POJZI6gBvKxydL2VNzm8LfV8GIYy6d+gseqFxRG5wgPLd1J6yfCKuCwWWopuPJRchAAnxnO55m8pkWlSwONp4o1LAYA8ge2fkOOyWef5dkl9Aw72eUIgja33GOH+0Fh3U7RUUUHzR+5MeVdO4FLX4TrIJF176FGlbCIUt0C++ZIatzMwMHU3M5hYKlBpMKWarVisKnqm48IMh9np5KhhsvFhXVY7asqm2I11GpqslN5b+3kHbqmwp1OC2UGVOolQmtU45csPGtfZ2BvfShfOt5qw62mO6eZOIz7EbTSCmXJ8wY5KO7aUMaBTbnaPsSJYv5PhcriEHV0mvs0pfpaqq2bam1FcW5UXd1U6aryZbNDTzdVWlqqoxujjm7KdQyFhBpqllKjCtmonnuri2K+bnSl3FQp5UFpCWWTtfrHMIS3F35A/yScEntSARgnUfiBlizml1F+4awQOP7L2onoTuh/D5mpVVviAm/tN51iDyonNlFZ054TvpCeHRVroUnWNIx+u2Cy3CeyDtpT8oc/oJmKexMqMKRxtQylbEo1zaN26Txqf5pHP82jFfPo2ePNowf7TZ52yeRpf5o8f9eT59kjTJ7stxwMy29UoYkSp6gH9Cb9rO5XkjJwQjcFORpiKTtloSzSllO6Z5QMt5zsjozJkbvCIq0eEcwkwbcmvMZlLbzv4LbsmAPzk6UbiVjwgQVg0okMb/mwJ4YT9jLWbQmWKhDc54zX1yC++IipZ3JHzmP8+T1KU4A8UIfbM2WKkqLbEpbhNhw5lllthaUH5+L1IpFi3o8WWUkXhzmRPP4PHXSvcAoRvKkhcUZvcs/wjYK0VaMoqO2YL+6bKtyl4Kvxm1z6OM6fcWep1SpWvGpxKYNhzJLdw7wI6oXeerGOVXh2+1OGw77JeHr98q9fpvQmheOuYeRp13jJiE8F34Vt+vk9ZEjgyLDmk3YUzBf9gyOd5YtF6nBpUGt5/P3tVQR6nixC+GpUkduHq8jtXiqy017VdURGqFSSYg0fpiQK4u9HSdJ71vT4aJMYw5xNwU4etJxWLlDmV2JiPF+6JyjfgoWhy8B7AQ6nV7oM6PnXtYZ6l6IfMKrSupy+osILvj99RllB4yt+6QJfTppd4JmtRtaki1FrEyNothK5ByxeDfwNuxmYefPp3cAl0EnkBDFuNv8x8i95ACAJVzAGzMvow8T1OgpXNEpu6zV/6VzSVkRR+/3gEq94YXtuQEherbEHgVYrvYC09WsYLpGAvSci3jbWYhchx2CZMMzealNLxc2bQxX1O9dZuPVrJ6pr5aLV/MWdvEy8XW3SU4Eypawl9iPFwUtopdfTsoMEKDMwD2qpsnD8vRtJBlfkxO5JRuGYWK0VF591oYs7fr+DsrJaf3FXGPi37ALHdp8ucZyF+qYVLqH3U7iSyG32IlfYxbBYfOfM6ELe3CFWlVj6a+3og5hbcP19Cf2CwbCQGls6X+BXZf2cpaQahJdfvcWOhNJfSfE0GYrflPyc3baNcFbbhhpIdyYbsVCH+dVH0lBUY6LBa9tstixbt9oDHFyLjdhAjeMZapdopMdRtr/9niFWH+1UdfWoaZKE3CPSK5uNIlvI6ljeCYXXIvU/vb1ES3FCtmV9Cfr2xZ2P+sfUz4wnxhOptrs40U3XPEJpYA6joC/4IJ86ESzDxArCPsMb2FFm0ohAjIoutDzzaAtTlpqf9TC+u0jO3P7me6n+h1VKGTWjMAHyOGy2xpZHL2tG2oZxWWhUhMO+sZhImT/21jiTprbIwzX/3hpcNKwrVdignKpM9tDOA9WCBlX97sEqkaPv7F0S6B4qVICe5aN/OVCqBVUwMPana0M1bO9ySEV1M6VshUwrUYkVrdRUWe79OeWHDhEGCtJkzGHuGNCxVI8mnrTCJEZ72yjbAcgCL8/DDY3LLfiP8hHyAtruwonj7/w4aYPFApZuMA9RlOIZhsxyuk94SIDxWw4rt4kKL4jvKn4rmNIUX1tJkFjfp9r4LElNusVZK0vQqL/7z+CLu7yPbM/faRt48AGMAmtlhbKXM6ReyZFFFJXU+GMaNS033YpTlJIA4HvEyrOdjSpuWwOIE7qS1zbMouDi5Ci7pFabkpomJ643D5QTR87kxN5FqWmZ5WISAGViSrNLxSQAdDFlyUUPC+3VnuVe4e0qsZgg8hmiSr6caFF2LgZOIecFjDf4UAl7roJeJ4rBXpAFALShqpc04Ul56EPTrXLAA3WlunzbsrhiQqekdG8pLlQS/9ebHMaer/HiBhyIKAyXqW40TULR+9VDkIuDLkZxFkDhEYdcQTEff9KUNttIV8vsjyU/EFcrnF3Hq75hqMWD/HipTbvPztTbeAmNvEKg1uDHlZjGFf6zsvdhPWQk5IFTJZ5pVHZvuHxeBs2p4kCw5zCQIxsHAimb9/SabchhnbzWMWSJ7l1zgttiZlorw4EqmQRNniVJ5M9gMq/XWAs35abVAnHzEF/aerwAn6BonLk1kGqnPoVKR/P/HeLjWS2diMjOBWDOrx0xf56dTJ6HjWIcI3umqbSDFY2wvQSSEuYS+QEGKqxnmlqb6DxoE/nj8YCDUZEHvA9LA7pfuK7ofmiE9o7WFSkJCyx7xOkR49ApzTITMwNQw3PSNeub79m7bjBxaCMdBoBw4xSL9WN7wTz0HJ/08oPLFwsfmHmjHAqWlnwCcIvepOL64i6lJFyVVkaaxV5AUO+MK0JpDECy+Ut2gmQmbhqmk3DUJR4B0mY+guZAZVwb4OVYq9lFQmnLB7WxFszOFTLUSO0QJl/CyfBT7TvLRcdTtMu3drGr2fkq9zxQUWdBigZhWRh8r7esdh/tL7vdoUvtkoiPqSFR9UTEbwvKkgtBaAsenHKv5Hv2DHtRHkzaUgMn+q4SRc/n60Va+ZR2lSbn/u0uPRZEMwTczrNTnhI36qmpeXKcv9yRQrEqN9PC2DflQY4kXJmwIDlHgi8KDpPrcX56NsNiGTke+6pgpq9KFlF5To7Lv0vIW8P7LFXLDpIesNBviygikYejqgWHQriLP4jJInBZAbi0ItGHr0p8q4S+NHSqEpSHzkisBu0ilq0gyNREEEcix1NyelvtVDpOTEANzRe2ZSjO97VwrV2vYJqhmP9NFC7fCsvUfIxQ2R2At514rzf8+gyxHs5S6zd+4IU3bY9eg4fBSmVAyuX8JTD42pWyuUcvxiZP5aLw667i1Pd/VhnAM+/9Ok6WPAaoFVRKtfzFHTeOX29+Ym4fe6Qkivl5/rpuQytykLG0uqsEbb06iK2gP9Eq0SBfSk9qnBYuOPm4JV3zFZj3oumsk7BW0kxszRL1UK6hvAntibElyzrVjgVn6GOsQOniKtnCUpbQOH83TuJeSXqr9xixKYV6vnrRsFCrJaJ/j7n1d/WIsoAye6X0i7sSxdt6KxyCtLDIQrj1hfIJL1vcGqOFANwr5jQ2M3RpNjKNB/JspVVcilukH3jBxcBDWanbghNZHrRQK1kxuGUXyC08vs4skee3OUFyClpd5wOl5WI3CheL54pNdscmXVMJ+Jwh46BJZvTKufbxIdcaXrHkBEltq5chO9KsnFdBEv7Rpzf1OwM60FyE7gdICagDSpQT3KqH3ovyXIbrmGkGylSPoaWXH+E+U5CcvOEUQ1u8cn9ukvzLL8qYlmIaTkBn7ijvluKxXhb4DQMq3SCpeXFlgGWXQWenIS+utWsSs81HWc32KZQJupw9vhatWOk73E22aZYdiC5ZQyuUkWjWy+4i2IGP3UVgc6gLYuabJ5WyywVccWVOY7KHIuiiloITYBoBxEDPQ4uJeulCWUaZqyZPPsoImJG/TaMEGrZIN6GzK8Xy7gA9XebjjFjkq68UkcnARxowd70kjlUHUZOW2PpQU+671WCEd6x5jvstHavVSquu6Y7BeTOyaeRgr6KPyoouq6wikW1JS/2StdS3qXNe0VS/ZE2VQctt9a3Bm9d4Y/22uq3S9eKPbKxfPqaxisPLPdrql3u0Vb44Xis9ArLX9BVyi0CfvfacYvaaYHYywuJ0u+bQN+lRgwfOpCl+cTotnopBgIdPjhf5rMgLneg3ql5cF2+ArebCfLNHFgBVZ57aC8yg3jEGyzVt3HEfsl6ZwnUnpp/9JiuivphlPlKknqjaJZOSg1gFNP1dAPMNIXvp6s0VpQuTrqbjo7NIoNxiYMuji8T5BV0Mjtto85ScK8QW18u/pHMHFEcPieZvUReM9ga+vZDnc9rauzeVlr4skyIg7mvX2BWuAcwaCPcd8L3Y02XP7qeTERVCE/0gdJGjM9LbVbvWKenp+qeUelLCLt5mU7jORmVX+XqknmA0HfUrrccJae2syNGuipyVVcQP7lWR1u6KGBYzZRL7+8DFDw9beik6lXd6JJd1TJfh/iLH/TCES55qOvxEG/PqmoQUaCWOy4PNUmF/1pd7HrM0ky+bt5fhutqVE8f+NT3mD7hslTUx0y6KqgY0BjAKF9byQNWPAX0pnu15vOtq1duF5ItGDZtfefBCAKUpBrgsXCyD8sSJej1T8Qq9/S/RK1yjp16CV7w2r5DPWUj03WLlF+MV70ViD0rwdxvE5I93TYCW4JaIY+mhJdAcMetqJ2Ozg0UP2EVaeoTHdp1VrQwqP71TCbZkEUfe5jWrPezTZSmsZLH6Ad7e0VI958K6gvmoiRkYjeP5gh0Grl1zlS6FrT5tZJBY6TkKHRJrlXnmtU67X9srvrsrfF1W53WCBbJGWm20IxsHVWfDskNeSieo1A8ZpEw5ZJiSDcgH2qLzTzgtFK4xx5yvFwVffE+NL6t1SlY2jkSSXveOvHe5AGY8/EIwKI9dAFzNbDnTZvu8YCZKK5veNLCly3el9Cs3bxeg0+BKq1/Fc+7F2nYVnLL93aWBfLFOsWhF320ToOou6RPCVtcHPgP8flSC1/cfphUi1vMPV4l04le1omzLUtaWxTv8xRnfO/FYgGYTHReTuB1nuE7/2JSI0NvJwbbBrLDkymcmwJswTH4IPVpvtK/COAFvcx7E7TQAld7ICB8nJ0/BcvZXydnJ0ySiNHZhCmhF66B1RSN6doIb3AnbYXV6OA8XHr7mcREA5cOzEyansxO2sETQcjg9dK+o+wHqcGjEuUjCy8sFoj5lSCp5dtvHhTObgRd8ePbVIpnI2bUgTFhm7ey9s2kzIRBgHiAUMuijBzlgCiIuj693m40SquTvf/svq22R//d/8ff53//2f9gjv3//23/D3yn5lUbhcXdHgWl2+kdISK4lg6YbSPGod6jm4RlxUELvwkWlAsWozIdM3Pkp8tIGvchUUUPWXtY4PHuTzu5cGdrtdsq3URGYlohmBtVf+C7rBk9DN6FJKwYcZ3l4lj2XxbsS00f+DdVxkoVD5uyJgEp1nXCwNvDwNgyD9EkNQSJZiYADjrTKQyL1WkKXKzRVkA6NIhAYGIV42CQ3X/mLCP/+9scf2sx5riPBtjgTrtPjda81lHECqCFK9qYIUftY2/QeBmc3f0wDx4kdHTCtS9qMSpUOK9rlfRyC/twd4n7Iw2Ny+JKH3LJLdcBjwEAh+MjEYYv/jTb5Fo3Wpzwdt3HxlxBYkPGwSQ6l9w+Q4jPtJwNhTysghOh1mKE9Y8Dw/9fXzzFLujMBku8OHQv+2OB6HOZHZSHl/O4QTwkgJkBALutd7CtAEeuYdDGVHRGAVOi7h6Dh8Km7ncLnW4UWpmgnjNJk9dYhxlFan/x+IiwJ5k4sB0eJw6WzyT6L4CZ8517VYXYVdAaS3QidpaTPISMX5+MB34XfZe+02t0e/9qz2dfBEHIG+D/7OoKPQ3xxln8djwDYwqfx+HO2Fn7vjuGXzV+vtTFhiO8v9nlCBxI6VhcSxvyd2B6eB8ASOl0OYXXwRAD7xQsZIPoYn+uzLJZg47mBocVSMWHYR5axpBHD6OLHTjf7jsTGvQwcSxxgpQfsJdouUOshRdueTlERlB3zIKIh16iywDPXlrTVxGN7hykOUxnU00MMjFSCnasg7BIPpuagNLki5s9nobalCany7SpjupuN/wwIjFYqHM/cTrfyACIGC/gIU0NhOirMXdqcWzVR/ZW8ZNtpjsmL1z8TSxD76jKZmEvhDOAExawwGLi4s31McHHokPgeGgfRRTaeeYN+dzB0HHvc6fV6g4EzGw574/HIG1sD0LKOYG+XzZAZB/Vzq433reP/oFPs9ZVTMSA1THzfl+GRS4fzft/udmdOb97vjOzBoEut+cgCrme2c2+GOymX7EGaJrmhzocLlvAT3gT6GDz3epbXnQ9HI9of9ix3PJv1Xa/b6TjezJ65lmvk+V/B7sPjTdWWX3/ETL+uMP1wQGO233g0RWky+2/wyf77ZP99sv+E/fbTanX7+lfna/p+5L76S3j08+s/X79+H334+qejS/r2w9FD7MFSm9BGm7BJuHHYrTQO7aJxaN/LONxpaHbMtDrbxzMrWzDyDEf94XgEo7U1tAbDsWRmGrJ2WJ32oN+3ulanb/V746496BWsUAPNT0bpv4hR2lGN0s7OMppkVyH2J8v3kSxft9fpWqPxoD8YWD0wKkeWY407drfvzqntOcPcKEN6Z/wBR5jnlxhXvlC4PQarpt5WjK/jKyeun3FLpG0yyxjOMbPgqIfXYDXYJLvAxeXPHfYDExXUjpX+21iGRBS/j4VoyDLKwtSqAL5wVjFmMkPToBgajc+NxuY/o1FaKYxybf6NbVjhq4GzNu4NB4MhIc1sLhqzz13wT8YwbE+zDQ4A3BuMBoNRb4AA/cF4NBoNGbCY5rrTgstXUr2P67rD2WhszTx71gVnddgbjKk1HM2Hvd7I6Xb6Tu9fzp9aM9+0blvlcXRU1Q7T006upLb1yXX65Dr9k7pOb6/CG2xsfPP8kkaEbY+JiROTgMb4tC9/HCa+dwD9EsPmr/DXCxZB3+v76WnRsWK9codbxf/n2I8WcrctsytkW48WdE83I+W+EW5Jypwh8SV3f9gGN83/EUCywyOSZBfnn82qty3V4gahfzK595u3R7NRf2R5nY4Nc3V3PrOt7qg3G1q0Y9mea90/2NxhVyiZ/2W2CJ89HyWS61IMlLsza+bOejPqgBliOz3bGYClZA1t+ttaHr8ve/0DvcXS59EttlyV8fMvYIqz6qUKq9a8QXB6A62EMfD44PzcHo6Hnf54YI1JB8x1u9MdDwbTA3JOxp3ReGSPrS7p9vrdwWhoDcaY0bN7A6vTs60+6Qzt0ajXtSyW0ekNwHy2LcQY9aweGPfDEcvoQEcdAw2bkF63O+rCTw8z7H4HsEHdOwQIDYfj/sDuYkZ3CB1jjPRIZ2yNLbvXH/YZqaE9GHUG9mhMIAdKhIK6vIxxZ9DvjoYdQsZdG0cGnkFGna7dG9r9LhQOfkUP42LT6f6uxEebg/wSvFceu4E7t2BmnovLTWOQCPR3kKTb7c17fY/ajuWNhj02x++9siYD71zVkoF3LifJwDvDHDLwTsdKAd41mits7Bo4pxNhaYduBMa10dZOzWFsXwEkwDUD+LCdqcAFZh8WUF8sQrZjNv3edjGBFcMeAJzkORGFKcGlf/KTq7qCnhGdR86luI5D24b4Unz9RkBgHVJoZStbShjzDTf5ZAqp3eTDTvUk6VUbGWX4KjZAPr99BcRTbHxja8Jw4nAdufQlu8OuRIaf48x3SI6Ihi5uiJTFklFrz/0oTstmNQOEPDf3PbZqSxiFvMvV+P8j6ocz</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_fb24928ebcf34fd283295e04e5f0c2f5\")) .filter((elt) => !elt.dataset['step1']) )[0]; root.dataset['step1'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "class Weights(nnx.Module):\n",
    "  def __init__(self, kernel, bias, count, seed):\n",
    "    self.kernel, self.bias = nnx.Param(kernel), nnx.Param(bias)\n",
    "    self.count = Count(count)\n",
    "    self.rngs = nnx.Rngs(noise=seed)\n",
    "\n",
    "weights = Weights(\n",
    "  kernel=random.uniform(random.key(0), (2, 3)),\n",
    "  bias=jnp.zeros((3,)),\n",
    "  count=jnp.array(0),\n",
    "  seed=random.split(random.key(0), num=10),\n",
    ")\n",
    "x = random.normal(random.key(1), (10, 2))\n",
    "\n",
    "def noisy_vector_dot(weights: Weights, x: jax.Array):\n",
    "  assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\n",
    "  assert x.ndim == 1, 'Batch dimensions not allowed'\n",
    "  weights.count += 1\n",
    "  y = x @ weights.kernel + weights.bias\n",
    "  return y + random.normal(weights.rngs.noise(), y.shape)\n",
    "\n",
    "state_axes = nnx.StateAxes({nnx.RngState: 0, (nnx.Param, Count): None})\n",
    "y1 = nnx.vmap(noisy_vector_dot, in_axes=(state_axes, 0))(weights, x)\n",
    "y2 = nnx.vmap(noisy_vector_dot, in_axes=(state_axes, 0))(weights, x)\n",
    "\n",
    "print(jnp.allclose(y1, y2))\n",
    "nnx.display(weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f26b99f",
   "metadata": {},
   "source": [
    "Because `Rngs`'s state is updated in place and automatically propagated by `nnx.vmap`, we will get a different result every time that `noisy_vector_dot` is called.\n",
    "\n",
    "In the example above, you manually split the random state during construction. This is fine, as it makes the intention clear, but it also doesn't let you use `Rngs` outside of `nnx.vmap` because its state is always split. To solve this, you can pass an unsplit seed and use the `nnx.split_rngs` decorator before `nnx.vmap` to split the `RngState` right before each call to the function, and then \"lower\" it back so that it becomes usable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8c9e5858",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "False\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<script> (()=>{ if (customElements.get('treescope-container') === undefined) { class TreescopeContainer extends HTMLElement { constructor() { super(); this.attachShadow({mode: \"open\"}); this.defns = {}; this.state = {}; } } customElements.define(\"treescope-container\", TreescopeContainer); } if (customElements.get('treescope-run-here') === undefined) { class RunHere extends HTMLElement { constructor() { super() } connectedCallback() { const run = child => { const fn = new Function(child.textContent); child.textContent = \"\"; fn.call(this); this.remove(); }; const child = this.querySelector(\"script\"); if (child) { run(child); } else { new MutationObserver(()=>{ run(this.querySelector(\"script\")); }).observe(this, {childList: true}); } } } customElements.define(\"treescope-run-here\", RunHere); } })(); </script> <treescope-container class=\"treescope_out_00734f5acfb64509aabdd5487c15c3ae\" style=\"display:block\"></treescope-container> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_00734f5acfb64509aabdd5487c15c3ae\")) .filter((elt) => !elt.dataset.setup) )[0]; root.dataset.setup = 1; const msg = document.createElement(\"span\"); msg.style = \"color: #cccccc; font-family: monospace;\"; msg.textContent = \"(Loading...)\"; root.state.loadingMsg = msg; root.shadowRoot.appendChild(msg); root.state.chain = new Promise((resolve, reject) => { const observer = new IntersectionObserver((entries) => { for (const entry of entries) { if (entry.isIntersecting) { resolve(); observer.disconnect(); return; } } }, {rootMargin: \"1000px\"}); window.setTimeout(() => { observer.observe(root); }, 0); }); root.state.deferring = false; const _insertNode = (node) => { for (let oldScript of node.querySelectorAll(\"script\")) { let newScript = document.createElement(\"script\"); newScript.type = oldScript.type; newScript.textContent = oldScript.textContent; oldScript.parentNode.replaceChild(newScript, oldScript); } if (root.state.loadingMsg) { root.state.loadingMsg.remove(); root.state.loadingMsg = null; } root.shadowRoot.appendChild(node); }; root.defns.insertContent = ((contentNode, compressed) => { if (compressed) { root.state.deferring = true; } if (root.state.deferring) { root.state.chain = (async () => { await root.state.chain; if (compressed) { const encoded = contentNode.textContent; const blob = new Blob([ Uint8Array.from(atob(encoded), (m) => m.codePointAt(0)) ]); const reader = blob.stream().pipeThrough( new DecompressionStream(\"deflate\") ).pipeThrough( new TextDecoderStream(\"utf-8\") ).getReader(); const parts = []; while (true) { const step = await reader.read(); if (step.done) { break; } parts.push(step.value); } const tpl = document.createElement('template'); tpl.innerHTML = parts.join(\"\"); _insertNode(tpl.content); } else { _insertNode(contentNode.content); } })(); } else { _insertNode(contentNode.content); } }); </script></treescope-run-here><div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtXAtX2zq2/is66bpDMhATO86TwroOTYC20EJoaXtnFiPbsq3GsY2tJIRZ/Pe7JTsPJ06Ac6C0PcBaJUhb0n7vb0ttX0ds7JI9iYWERIYfkMvQ9xn6Lwr8iDLqe00UEhczOiQ7yPI9VrRwn7rjJur7nh8F2IDxkUMZKYpfmigIYcSlESuKrYtsHMCo53swrGOjZ4f+wDOLhu/6YTNeuoOS33QXCGA/ajKniSzKgMxjxGM7qE+9YjIul0r/A3v518WI3lDPhnV+aJKwCEM7KMCmCYNFl1isiRTD4dx4pOgQajswIksVfp7HMAXhpvsnH4pDGlGdupSBiHjA/CltkXospF5EDX4siWcTuW5fb8d6fD3VYzEceHBmCGOREdKAIa6I3Q0cBC41MFfttm8wwtUUEtzf2MvnC7t7oHk4L2LIJJYXoV3EHBpJNmFnYJYT3yT5guT4EZPEPIhGGLoMiMdF1gy+K1/0f//OmjnEnukSmPYGrrsTnyABm13f92A0P/LDXgHN8+BfwBCfSg0zavDBgISWH/axZxDJ80f5gnAEOCC/NIOK8aLXqKwUYB9qofwC15JLPJs5aHcXlTjJWtZDwgahB3pHxI3IjDFn4HHOFreOHGoxzp8g4B9u4XvFCXlwP8/0R1JIrgYkYppH+8JcnRD3ST7WSYHvsbN0UDCInFiNOxkyTo7YjcVYI+X9eeBcxIZkvm27cfheihADbw34XnyEuGwLkSE4eGJJzp34XeqRMVd6LsxxhhJiyXBxFL2HKE72zeeme172wQ1zk8NvC6BPcH/h43uvt7MCwKRDJDbczaXzTA4xrIOk5Ho3V8pB6IZsmcT3gEVQhgdT64IhWwN5vmYiew6CMc53us+Y3+eJoen5LC9ZvmtiHVZ7sG3TwVF+z8U6cffSM5fxGWJN03CI0SNmoYD+yVU3STzMD5qoJMkV0l9OPXzspihEht93MtPsrSTS4SXW9ZAMhXeL7PiqWldwqTQjMPx+H8Sao8Dii5tmgQTHLDv+kISFDPqEPOIpzp7fsCxX5MqUgJiXXKIZQZKxlxP7EIf5YlF3faMXDxV2JmlaKEgOrlHku9RcRxlb6S7iW25/EnLmAmCHOGAxEs4LUePfSf0SDtBElGHIw3xxysArKh/oM8sPgNqkERw6nlS4RUK0h4QfNZs6gZw4p7lXhvjKdoG4eBVlXr2Sqgd2n55FPeFXQg1rzhTWXj7ZxGEvItiGOPOWVz9SPEx54EuzF03oUxyKKtxEG/9SKrqx8ZzspRetZLL6A5jkduQHD8KIGzDwAYuQMONcGj3esSIUxEFFkT2jVT7+OKfOxGPkmi2fItHo0qJhxC59L05Cy6G1LpQkpcKjKdNU6C+zH1t8kUUuVR+HNqDHmA0R0Ld/8TRIxcFYH0Bq9DIT0Gw6y2lzKLdABYoE7J5N/C8iq2ZuAfpvHGPwCopd1B33dd+N0IcB4/KaaD9eCT+DMQRGcUT0HsD4OPP2oTI7ArBjj8FyiiNiTsH/K1Li3zvLbh6vFqC7JDV4CU1LGcdHhhTZ6W62Uhrh6NKAOgCKna7HFktVj0meXnfmwpr0kfOqTwqYiRkuYg8MK2BdYX6YH8Kxaoi9iTeLbZEcIQIaA/RQ9AfsYaJMOQDDUGL+keZEHIn+oP3ADxn2lvbWQ79HvEs+MktGd2t3btmcPidmvpUcw7kUXeMlj445CBUHS0mKI3aOLuRAao4wjIHVlJIDKxDVvDSgEzFD4iXCp3s94DJN+GgYcBqMSWlOBb+BXSMPDSl0QwBoBICWIob5+im/T8ZJgqhiTkzAVeAmHFbNmeNqgF0P+otLaOEteg2bpAKvLgIPminMkdYIhx5Y4HJSKibWtSxsyOUMwgC6kSXDiXyYKCkZKibGnEHKpkDROCzaITYpmC2P5HLFJPYW8iFIbIJKwF7VcLbioIE2gqcgMYQSNS/xspSrH6cQoKWUX5w656KtV/nmIt0Uesc2hOhOk/xw7y1JquEU7sPrj3HnuDUQXGU2CKnOTMgBNnEdsFZeBnHmlyT7RkBFCrzDTnQwW1q8/9rJuZM2RqycbjkZLSbDP1Cbs0ZMaLDI42oQJQLKpA/SrRYsA0M9GMHc5eL/TfWKsX15trpHr7ie+FaCHZ+k7Z8Z+a90/iMHCiZv/LOQHQfaoLxMmrglnKMCXldtFO8AKnFxAAjs7i724QZefcKM0ZiIXANUMVfSPAYfWUekVDG7L8tSRfpCSkpfhKE1OyyK6qWu5rIJ73PYyn0e8TaaXxCiP7QwxGPJCv1+3vSNAb9KkjgoiKQhdgcEHK0gRX6f5AVU4FeL/KcUNwr8WvGerUJuAwpnYXqRGzmEMH7bS0Zov9vtcmm6fIzf3YpJKSTirqc79oz8f/43aU8MMgEtD29V5i+HPH577SZjoySGVX71FoVGEw1CN89xc5PPb498y1J2dEDmVXXLLDUOjm2tpYmvo1NN88Wn1tkI/jzsaFpbW/fV6mua3fPfmUft1v7oq6adf91/qx0ftfa1jn19dPjeYVHrmBK73HnzRXl/VP067AYD+vG4ci6//XJ09vl4eHF8wz6OO539zQu7d05bb0oOfXM6eNs2D76XDvVta3hkBlfvqs7VBaWng2PvwDm0PjHtU7V1Eqpa58jrtavGp8HA2zyrXBlRbzS0Ou721bXd9uu2/nZ0UJcPtW1PO6u8D8O38tmmfVM6M0vaW0u2T2r7o4Pvil3yx4OzWq3flqujwy+ND7YdkPPeWCVH+k3F0MMPBwxr9unRyegNjsbR6eDo6MtFuzPSPp4GR1/NT9vbm3btvPalzErWu49X2rACe77XTmra8Ujr2zdn3c3Bty5pf7lWrKpxc6KeHY4rg5b27qb1PegEZXp4ut8ufRt8VLs1z2q9bx92jvsa3awP24rjyU5tU/88+vJ9dBgO3xx82ve+W+22zTY/GN9ct1Zp7L8dtepOQz0+PuiWD75pdv+o8r112mDnB+Sw0W61jg7Kb2z1bPurMda1A7Dp53fb2ukB1sjxvqsd3rQ/2N+YXW19tD98OHrT6tHTCum0vuy3OgYtBU7oBx74RvCt/Ua+kXtda99izvidd2jiTnRolU76B+2TasvUrj5/DjCLut/6polpQ7FuGuon+v2qGvTD6gf/636Xhgf94duDcveiW+60FaN1ap1vHrp+cKB2olEF21fVOv1GuiducOG1Do+IeRySwcXVwX5fvuiEvW73uqJULy6ikQYcFZB4uWH5DeHWG7xk/gf+mEY/Nv0AeodZSIr3JkmS1lBsxTH7b9hr/Q2+Ix5ARMMY97KwN7iHZ6B83FKmn6cgBM99Hr5AlrScfCyC9MC34D00bzzxCFOGPDykNmZ+KMHOge7j0JRGIWXknFyz/GwvjijivWZvIFDi87m5Bpu/fsAp57RPoBPPT57HltaFpA/d8tLS2y2klEolgaQg+QJqzYurouxz57ro3Iw5fkk2yWD8wSiHXqEOpi4kNuYjTvyHyGyANj3o6yAbU9AZwSa/BNic113yknPHGw6/UEAiPe7mUqCqifyeazgAqWv1Ckc/pTKSazVJadTK1bJSl5UK2gYABOJm4Ul+sZ1Ltk/eiNKX2Yt9GhDHAOA19YJBUshyouTr/nUuc5MEHcBkjAxARrE4fW7qhSOBiSh9R7HAabr05/b+4TKuR6BYT5eaXGi7YTbZIf5xIcpPlJ+OJtIv7T9BObkMqSbvQLk9cJOPOMQAVBsoX66iVmEL/GTAL4dklFfF72ee3WVQN5tIgS5HEWPnPsMu0MBIhY+kmZw/cWOCPDfSw3zo0h266eHcUlecbgVya6iXSHWKo93VvvRXXCfbZZ7GASwXX0uedw0gC6ALcOJSXUorXNgwn+UMG7MLuA3ke/s87exuPDDvilfTwgaa3g7u5iSu3NwsBUynAGSJOX7du3ixCXMix0ANc+BzIvjeaud5oCuXE/98NG9cukDMrZvP7QkYvHsPW1NzNzd9pmyo1boKxVyXia4aRMa4UVWVSqkmlxpErhsLh2Y9bya+li3pd/AfAd1Rlkam/aMgnJBYro9ZWcmXtwr/sGdefIelUq1IpqyWIddrumLJdUtVVblUt6r1MqmbRsMySrpO1iX/xwnYrIz586j07hjYemicrMs2T5oe4q5wXZKIKR4nVSymA8gHn2kEyZTeiB2RQ02IVgA+aNYt8zb6Tzn1b6G++/9Y4d3rDLCq+M3/bZbcHrQK9zDmcrwU1hfeJfqtZfrpfRQ4ywNj6hesqqudfD2IMjga/A1QVFoBAuP+WMAkFLnKtmLyOSBTgvJ/CsSUAgo1VTUMpaEaDd1UZaXaqGJFV+uGrFoVUlGs5wBF1OP1O6t6P2K9floHXF9U5kietKq8lJNfP+f82YLSI6FH3Je+/IlMHqt3lc3j2ecoNFWUV9RfrzcvyzVZUeRG3ZSJSspmQ5dNpYJ1rNZrFaXxrL25soXKj9udV0sylsvVWqlWN9R6tYqxVcaqaaqqZZVlxfwbdOd3KPV36c/jRLAeDczTvPToz6jCF2D1swGrp62yfxZahZ79VA8e69/KnhhggVzL2OoMhP2x0Iqrd5XJ+dxzwKrlx7lnA1ieTyPyd3K/rvjLUT/eByWh6XWeGFP8zf3xd7m5vNMPn+Eyc+Zk6+8YFulerjjn2pmKYZXKMnzhhq6WK7KOiWnWDLVWbzRUQ5Gfo7ccrLzj/EWQ2aLHrcfn2dRPgdwWvYJhezfDceN/cg1mfXWt1HYEZ/HHX+mmeUmrIOy9LQC0L3fPv0GL9ONLxP2Dr0fGfwdg8I6Mnw0WgIrvYXGgeiZIUP8pIUG1bCpWqVIzZL2qEqLqZVLDKjEq9ZJsVQ38HJAAbMR3tMIxz0W/BzQAme4NDKa0L7DgUfV/P1CQUL5Agt8LEvyA8vASdAs6v1/AvQTbHcH2i5n9Ke8JX8z9E5n7aZ4mfkUTr/53XLP/6A8t/DcyuYc7xv27kQmlSYd7/w+Pq/nE</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_00734f5acfb64509aabdd5487c15c3ae\")) .filter((elt) => !elt.dataset['step0']) )[0]; root.dataset['step0'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtfet62ziS6H8/Bdt9oRRLiqi7Ldvz5dK9nT19ySbdM9Pr1edQJGQzoUgNSTlye/T/zHvsPsB5hfMo8yRbhQsJgCAlO+mdb6fjTtsSUFUoFApAVeF26gc3VprdhuTs0A/SVejenlhRHJFDK/DPDhdxcumTBUkS4l8uPGcynvcWzmQxGAyc7mQxmvTJxPeOF153PieH56fpyo3gN9I777hJ4t7eBL9eenGUuUFEEuvOen8dZKQNcB7BgpKlG06trWUC7gTRIgaUBaS0F+4yCIG3ZRzFFHtqeXEYJyfW5y79mVpLN7kKovY8zrJ4eWJ1O70hWU7VElcJqS8uiFbr7CK7XYFAEje6IoczYOGGJFnguWHbDYOrCLgIfD8ESosgzAjwcAXUUsgnDadpxVBUkN02up1h896FnVzHN1RQZdL3oxetl3OSAMEozhoni9hbp00gO48TnyTtxPWDdXpi9VebDyPJPlOmkbxokzH9mfLiTixntbHSOAz8Iqum1E4KkCRJdX2paz3KQhasAEdR5Km1itMgC2JoNncOPKwzSJu73rurJF5HfpuzTAsyMTwPARaouL4fRFdMr7xrJBtE0EJtckOiLBWFvQ/87PoEWi9rI3OQNbWQs0UYvz+xboI0mKPilKv1azuIfLKBkrvdbn0t5/Fmz1rGm3Z67fpYdJf+h9WiFWrxhB4k8KqbK5TzdVzDlhcG3jvfzdz7tFgYuyjRyyVJU/eKSNojevT29DEbS06zhJDUi1eknayj9jVJIC31kmCVWVQ3bXe1Ah5clMDj2MtI1k4Bx13a5wf4A8WmmSW4sM6sRqNpnZ1bdweWBf8W68hDVMsnKUkC6OW/kp9BGpMGdhoAsKyEZOsksmjqE6TTWSTxsuFm8RyAWlZjSQkuO17sk5coyidZo9tsTgF7e1BdzDcghqzfKwpirM5vM5ICnw8qTxBZIG2kEpH3Fi+I0mpQ8p35egFjO0fhFWQ4u7h+Ef2P8EyLuRfHjOWQZNYzVKalu3r1L0+fg2ZO9dpckewZKGMQreN1SoEbN264Ji2mhoCJaKKGSHHupuSS9oaWFS8WKckYH8HCYqjW6ZnVFRiWBA/V6U55KsMsUrYWCVMiETk/s5wKIjJnnZBEV9m11bZ6JdJORyUuiDERe1maU2RFPrIaZtJOc2ri43s3u+6A3EFmObFmiYuinC8th/MjtXSiVeiiKGJ20Z0hUw6wwMg1rSNO3qpCso4shyPKrcMKu6orzHloYY65sHldYb2HFtbTC+P6f5G0rKuWNZ+ZO+1tBKaT9yTx0iC6fkWAeoOX947c0jH/j1ztw2D1vQu6nbjvvw8i9he/cxL/4q6EWubU08yF+ew12ic+L6IBGJm7LjQYNfuzIP0miGBiaNCsv/6VqRBMVY1N03qMCNap5ZD2oMDLK7gRmqVpcw5AaaVgmSGxR5TYoxwGfyhAGF81yqUecey/JBk0Cv+2it83NgygZfWazVy3t5IW531fkaN1powBmM/kqWVo4keNYfKvrKwOz8iaONMh81orGWIUpa0soJbupiHanTPUnFbU9DSHyLn8B/1h3Q5aC/uMdXagtTzZrBq5CqgygG5YqDfoTg4mqpY3fV4GiEqU9Jh1ZFUsjzXBsckDBAtjLeCaukze9bB3TCXVEVia6uTEoMc4XVUF8uEVu/qK4CiL2g7zbl1nfcxpClq8Zwkixtnkg1urUHohyeKrSVV1QihXn0TxMojAxkhyHQ6ihqQCpmprQx8XAWVBGu1aO6iIjoLoSrMpTGltpzJsbEAu/L2aTSJXTLGlWSCLX6Mh/TpLwOJmc71qu+WTiJi6lCnmTXI1b3xxl2ytL+6u8Nd823xjnG7Qvk/cFMyqq4eVKEGgOxKBV3MLEJ2e0xtB/0xgiO6MnWEPPl/h5+64h5/nxSBVoJ1bTm9SyJ5XxqYOkG1UaQFCnRhbEahWUfDPw5duBg5TBDYQtAf8u22Bm0iT8okSGrmBahpQgw/+nAoQbmRB2tFRU7PQkvg9wHPAi2AmFCQn95aRewvkADYn9VYmxSef+P3F25mcCoVkmw7y/4p4WQOti7fAO/wJWpbTkmy+QiO3JdVijPrBVZBR4/llEizdBJvqgsLany/oj92Cj85iPJ4P6MfFYrzoEvqx57ndnkc/+qPeuDehH48Ho/Hcpx8n3nA0mNstTpD0x2OvR3Pm3tzvsY/OeE68hQ0wVEw6X68JpPgqZ+MF/kexXeKNyYRzNp+POQ8TfzFxeerx5HhEP3rDedcfso+DY+94kHO2GM9HPmPHn/vzCWP/mPguGeacHeTceSQMX4MXBSyNpyxDc1rAM1kEVyWfxYcR58eIPAN8McJR3YNGpm5Ly+IuTJDCyBb4hS/DCLbyAV0oBDfOKLSkhVxBaBcGLu34XehdN4bDLzFo0LSnBwZFgqKgI44pN+wD/mtO62mOuxrNUsfidNEIzmmzLxe5nl50W1bxb9ZSMhya6pQzPgrGrFny41ShdzBCAdMAVNnLHU5b7/IeN/sV66VESkAps02Q/uD+wJzHptzRSxKXB777Np+BHE4KveGQygP+NiXKD2pEELEjidiZyfNvZbPkDeOUmrIeq6qsWVMzwQvXHNrwBbQDODK3iqBZC2JTgVwMgQXwx1tWqe3KIjVN0x+ruWpn5g9vOSZi5/7N87AsY1n1Ldd+UNN1fxdNV0jU1AhV3ae+P9Z3VWdHp9ObiE9cPjq+fA7jQ95UtXcE3DlGwTS756H99L7Nfd8Gf3CTP7jRazvYrkxH6n/OvTCrM2ey1Vn03rw5T63uvZuz+/trzm6d3LsPbs57kzU0p9x0hdkiGriptm3u49/f9jAWWCKh5OGPqgIGbRKcSiqlNN32wfqieVm7DEnLXrngLGQ8OmyXwm8fZFGqvUpkX9CxdiabnRwE+FkDGzC1Ev+T9fnxrU99IqyIm5hklIdtWjRu05ICN/dvgr2Vk/re87hSL9GJfekmWfr09jmC5o45lYvscKH8xrNyCjRbr2X1zTkg2OFOiBGHGODfPijQTsghh6QYg4dg9PHv8MMw4e8I9aXwPuUYcxDdoB8PAl240EpyZ2VW8JeWY32mxSML6ynHzpI12aGFEblys+CG5EuIp8UKp4BZuldgbq99ZQlCt9mkeA2u+uY4HRyP6VDc7KSrMMgatq2ZegxJLFaelhSL55iMBgS9XCEslKnjXSiEZ/IYj1JeJXSHxGVCVsTN0st4gavV6zBU5nFD4E8hO7WOjgJ9zuM9PM2Am5aVBj7hPHAuGctSQLAkQwD8gW7T4dIB2KYKnEuODt3lyug8lYcVQ3CNiUxjy9pPVBXzdV25PNh4n1IpqFas0T6gY2HJOihmEN44xSxSai35c6WFoBlDoj81662W8ixaMpWqJj9pXL/Qqndhtv0KJ8+YYTAIZxXVNTSwXH/jwCObRXKwiwF/sjQ+pqVh3Cuytxj2F2/FUgqaEWwp7knkv4j8wCNpQ49kBywdP6SgJxRV24AkBvQLPiTQQVisJQGSBSOBCVtoFuRc2GjH2DNqycCsQxI3tGVlo2V0Vuv0WiBQRm1j7KlMUrfYBev5phlezQuO6m6C1J7N1IlPAJ9ZHCp9F6wu6Thka0s9ErtvvrgzgG9P1GQS+ZD4xuyN84JP71suwzNR3YVjtU1l4c4kXK+q53uf9mCNpy+alRZEyuxRxHvUyG4s3fQd8a14nTXtB7F5Gcbxu/WqxK1Yv7G++sr6jOMGV1GcoINIR8ua1qnmq1wdpqrpep5mYKPRvpurIOPtkq5V2zPNXRScqqhVnmOJw3X0LorfRwp7FVaDhCcXVjUv7SN77IIm0aN1Z+62HUQx99nz3V0np/mQHlDWq3s1GytSY33fVtvZZvU9pKK9tlWOiFzc6R/Obc2piEPSIUkSJw37Z8aLPPbbfB4x7uziuwBYAW/jIBK+h7LB9Ak08usV8UqrtJdg/Lm3P0dZEP6Rbfhu+ARjf3R7cstyKZiQnTR/0e3hAcw6tz/OU5Lc0M07fB8sSVJC8URWowE2bhKQNN/ILNqLp190Z51AQnyFxYMCdsuzD9+K/b2bvMNt92eWxG/nL2uS3L4Ge9bL4uRJGDZsfeu2LHpWuYa8LCFcIRJih9EKU1UIQDoJWcY3pNE0qXJZQh0/SKESEdoeemO2rLttvlcYqpFmTyJwHJDBbxJ3SaRN4BXEY/ZBbj9hyph3ds/XQeg/4fvMvwmu1onW+B6Nloha71IVlcHLfamrLMqqeU/+hPG0iFkQGwM9+E3e7Up31T91UzIaFEBSYgn2OQsVKaA0TYak89b3MDXplLUMGQdPIGDQwKfb0Dm8lCjDbkT3lUCLNBny1gB5a4RMQ5gDfAO4liHjqFG0AsXTdoJIVm6wgZHwJUlwG0iBoCRrklyTb6iNjesPLyQDWJFqFdD0oBinsBvT9iq2W8lNCrPlgh0zkGZKCsDGFP3Qg6Ql5V0UGuEg2kmWnUqoJCows2v0TXBc/ZrND+soXa9WcZKBGeTTmb9Z3q5O9e4SjSW1UHZORNPKpiy0XHJxiE4u3wqPCd46SWC0VhNTsqJOTFf2YrRQUr5Xt1DZjgg+3OpJTW2PGRbA5k1efhH7z/lBVo/E93xnP+VfSdzq9VRrnMWZGz6Lw1Srdxz+CQ9K0Xo6syKDVQekWnLjNAGU6m3cR5cCAOAUwFLojMcxcUcu5HTAegKbin6kJpYAKzhiIitME7luxedHQPKIihjKwp3+kXKMIq85o5bjKQqXS+5V/F6THGjutyS4us5KorvdV3S39xHd7QeI7rZedLxyxecdoiuqLskOEZsVqsinHA/jaq9xBP6RWrdY6N22Wj7aYL2HkDQMSVLlwi+onCKwPmbSGaADnWvPBc8tlZfuGn7srZfQ8TpeQtyMfB0S/NawGaidH6OiXzv0ICLu/S5U88jq4fEIsftQAb+mks3haXuY4SWhPqO4uE5KNpnCK6fKV1Eht2H3fMqiYZDgO67p2ZT/wzdmK3HiImYP5idOb9iEqmkkrYKImfNVsQFeDjnXLV0Z9+cZUTjPqvXK9/lDgSX4pXQcQNrcX9odznf36zvDS+vLbE8MPQHyZJ6aSswz1Y0JOaK7qUGkmSWGzQ11xuRbXssoNQPf0E63HKNAcMs9bt7WhWNev+cGJN0A/70gWXO8qfgx8d3CUluFCFuFUFpWt+MMm9N9q6PwhIlHeIDvcXEub2cwnK+X4SEhg+oEkaEJ76Nn7PDRUj6os6tS3c5wjxapaeE2Voe2MHLKvu0IwhwYTneA8f5CP0EphXXpAAw+pTYCK+ctchJHZ7KHIA/Ej+oH6qnOIXchivlEMkdQA9nY1NX2VNwW8Lf18HwYK6Z+C49VhwRG5wQPLd1J6yfcKmCweWoluPJRchAAnxrOF7m8ZmWlE4HGM8Ua5gOA9Qe6fmOdWJ99VmRX0DPsZJcjCNrcco8d7geldTtFRxUdNH9kxpR/40YeeRavo0zWvYcaVdwiEroF9s2R1Li5gYOphZ1DwYRBpMJWarVisKnqW4wIMh/nZ5KhhsvFpXVY7asqm3I11GpqslN5b+/kHbqmwp1OC2UGVBoVQmvW41csPGtf52BvvateOt5qw62mO2e5OIz7EbTSSmXJ84Z1UN61oYwDm2q1fYgTRf2fGpXFIerorPBpKvW1Uls39dqK4tyourqp01Xly2aHnm7qtLRSRzdGHd1U6xgKCTXULKVmHbJRPfdWF8V83ehKualTyoPKEqoma/WPYQjvhEFE/sSdEmdaA5hmSfyOVCzmV1F+5q4QOP3L2k3ITuh/jampZS9xgdf+TafYg9qJjVfWtOeELaTnR8XaaJK1DKPfLpg895Gsg87M+sMf0EzFvQk1GNK4WoVSNaWa5lGnch51Ps2jn+bRmnn0/OPNowf7TZ5OxeTpfJo8f9eT5/lHmDzpbzkYVtyoQjIlTtGIyHvxWd2vJGXghG4KcjT5UrZgoSrSVlC6Z5QMt5zsjozJkbvSIq0eEcwlwbYmvMRlLbzv4LbqmAP1k6UbiWjwgQZgxESGt3w4U8MJexnrtgJLFQjuc8braxCff8TUc7kjFzH+4h6lGUAeqMPtuTJFSdFtCctwG44cy6y3wsTBuXQdZlLM+6NFVsTiMCNSxP+hg+4VTrE4b2pInNKb3jN8oyBt1SgKajvm8/umSncpBGr8ppA+jvPnzFlqt8sVr1tcymEos9buYZ4H9WJ/Ha5TFZ7e/pTj0G8ynl6/4uuXgt60dNw1TnztGi8Z8THnu7RNv7iHDAkcGdZ8REfBfN4/GNJ5sVikDpcGtZbH399eRaDnySKEr0YVuX24itzupSI77VVdR2SEWiUp1/BhSqIg/n6URNyzpsdHW5YxzNni7BRBy1ntAmVxJSbG86V7gootWBi6jPxn4HD6lcuAfnBjN9W7FIOIUpXW5fQVFVbw/elTygoaW/ETC3wFaXqBZ74aaUsXo9pTI2i+ErkHLF4N/A29GZh68+Ju4AroLHGjFDeb/5gEVywAkMUrGAMWVfRh4nqZxCuSZLcNO1i6V6SdENT+ILrCK17onhsQkm839yDQbosLSNu/xvESCTh7IuJtY216EXIKlgnFHKw2thA3aw5V1G88N/QaN27S0MpFq/mLO3mZeLvaiFOBMqW8JfYjxcAraInraelBApQZmAe2UBaGv3cjyeCKnOg9ySgcE6t2efFZFzq/4/c7KCuv9Rd3pYF/Sy9w7AzJEsdZqK+ocAW9n+KVRG6zF7nSLoYw/M6dk1De3MFXlWj6S+3oA59bcP19Cf2CwtCQGl06D/Grsn5OU4QG4eVXr7EjofRXUjxNhmI3JT+lt20jXLfjQA2kO5ONWKjD7OojaSiyqWjw2jaHLss2up0RDq7lRmyixrEMtUs0xXGU7W+/Z4jWRztVXT9qmiQh9whxZbNRZKGsjtWdkHstUv/T24u3FCPkdLtfgr59cReg/lH1M+Px8USq7S5OdNO1iFAamMMo6DM2yAsngmaYWEHYJ3gDO8pMGhEso6JzLc892tKUpebnPYztLpIzt7/5Xqr/YZVSRs0kzoA8Dpvt465PrmwjbcO4zDUqwWHfWEyizB97a5xJU9vWwzX/3hpcNqxrVdignKpM9tDOA9WCBlX97sEqUaDv7F0S6B4qVIKeF6N/NZDQgjoYGPvF2pCN7V0NqahurpTtmGolKrGilZoqy72/oPzQIcJAQZqMGcwdBTqR6tHCk1aYRGlvm1U7AGng5Wm8IWm1Bf9BPkJRQMcL3TT9LkizDlgsYOlGixhFyZ9hyC2n+4SHOBi75bB2myj3gtiu4tecKU3xtZUEifV9qo3PktjSLc5aWZxG481/RF/cFX1ke/FG28CDD2CUWKsqlL6cIfVKhsyjqJbNHtOwtVyxFacsJQ7A9ohVZ7sbVdyOBpBmZCWvbZhFwcTJUHZJzZ5ZtiYnpjcPlBNDzuVE30WxtcxqMXGAKjGJ7EoxcQBdTHly2cNCe3XQ9a7xdpWUTxDFDFEnX0a0LDsPA6eQ8wzGG3yohD5XQW4yxWAvyQIAOlDVK5KxpCL0oelWNeCBulJdvW2ZXzGhU1K6txQXqoj/600OY8/XeHEDDkQEhkuhGy2TUPR+9RDk8qCLUZwQKHzEIZdTLMYfkdKhG+ns3P5YsgNxdunsOl71DUMtHuTHS206Q3qm3sFLaOQVArUGP674NK7wn5e9D+sxJSEPnCrxXKPye8Pl8zJoTpUHgj2HgQLZOBBI2ayn244hh3Zyu2fI4t3bdqPbcqaoleFAlUyCZE+yLAnmMJk3bNrCLblptUDcIsaXtj5egI9TNM7cGki9Uy+gxGj+7zE+ntXWifDsQgDmfPuI+vP0ZPIibpbjGPkzTZUdrGyE7SUQQZhJ5AcYqLCeItWe6jxoE/nH4wEHozIPeB+WBnS/cF3Z/dAI7R2tK1PiFlj+iNNHjEMLmlUmZg6ghueka9Y339N33WDi0EY6DADhxika68f2gnnoKT7pFURXz8IAmHmlHAqWlnwicIteCXF9cScocVelnZOmsRcQ1BvjipCIAUg2f8VOkNzEFWE6CUdd4uEgHeojaA5UzrUBXo61ml0klLZ8UBtrQe1cLkON1A5hsiWcHF9o33khOpaiXb61i13Nzle5Z4GKBg1SNC2ahcH3RrvbGaL95XR6ZKldEvEhNbRUPeHx25KyFELg2oIHp7xr+Z49w16UB5PuqoETfVeJoueLdSgqL2jXaXLh3+7SY040R8DtPDvlKXGjnppaZCfFyx0Cila5JQqj35QHObJ4ZcKC5AIJvig4VK4nxenZHItmFHj0q4IpXpUso7KcApd9l5C3hvdZ6pYdJD2god+2pYhEHo7qFhxK4S72ICaNwOUF4NKKRB++KvGtCvrS0KlKUB46E74atItYvoIgU+NBHIkcSynobbVT6TgxATU0X+iWobTY18K0dr2CaYZg/jdJvHzNLVPzMUJldwDeduK/3LDrM/h6OE1tvA8iP37f8ckNeBi0VAqkXM5fAYOvXSmbe/RiHOuxXBR+3VWc+v7PKgd44r9dp9mSxQC1giqpVr+446Xpy81P1O2jj5QkKTvP39BtaEUOMpZWd5Wgo1cHsRX0R1olmtaX0pMaZ6ULTj5sSdd8Bea9aLrrLLYrmomuWaIeyjWUN6E9MrZkVafaseAMfYwWKF1cJVtYyhIa4++9m3nXkt7qPYZvSiF+oF40zNVqiejfY27jTSMhNKBMXyn94q5C8bb+CocgLSwScre+VL7Fyua3xmghAO+aOo2tHF2ajUzjgTxbaRWX4hbiAyu4HHioKnVbciKrgxZqJWsGt/wCudBn68wSeXabEyQL0Po6Hygtl3pJHIZPFZvsjk66phLwOUPKQcuak2v3JsCHXG28YsmNMnurlyE70rScF1EW/zEg7xt3BnSgGcbeO0iJiAtKVBDcqofey/JcxuuUagbKVI+hicuPcJ8pSE7ecIqhLVa5P7es4ssvypgmMA0noHN3lHVL/lgvDfzGEZFukNS8uCrAqsug89OQlzfaNYn55qO8ZvsUSgVdzR5bi1as9B3uJt00Sw9EV6yhlcrINOtldxH0wMfuIrA51AUx882TStnVAq65Mqc53UMRdFFLwQkwjQBipOehxUR8sVCWU2aqyZKPcgJm5G9FlEDD5ukmdHqlWNEdoKfLfJxbXeurrxSRycBHGjBzvSSOVQdRkxbf+mAr991qMNw71jzH/ZaO1WqJqmu6Y3DejGwaOdir6KOqoqsqq0hkW9FSv+Qt9a1wzmua6pe8qXJoua2+NXjzGm+039a3lVgv/sDG+uVDGqs8vNyjrX65R1sVi+N25RGQvaavmFkE+uy15xSz1wSzkxEap9s1h74SRw0eOJMK/PJ0Wj4VgwAPnxwvi1mRFTrVb1S9vCnfAFvPhflmjzwAqs489jPMIP4JBss1bdxxH7JemdJ1J6af/SYrS30xy3ykSD1RtUsmFQexSmj6uwDmG0L20tX314SEJl0V46MbZlBuObDlkzBzf0EXg+E2Oyyl4Aqx+fXyz8nCBcXRQ6LFW9Qlo72Jby8U+Yy29u5NraUvy6QMiPvaNXa5awCzBsJ9B3yHe7rs+f10MqJCaKofhC5zdG4NdtWufWYNdP1TSj2tYBdvsyldZ6Oyq3w9Uk8wmo76Vdbj1GrvrMjRroqcV1UkiO5VkfbuihgWM2US+/vA5Q8PW3opO5V3eiSXdkyP4v4ix/0whGs91nT4kTbmNTQJKdBKHJcFm6XC/qwv93zM0ky+bNFehutqV26aBjfkhD3gslXWxEy7KOoa0BjAKF1YywJVP0bkOX+25+NdV6veLiRfNGrY/MqCFxxIpBjg8nCxDMoSp+r1TOUr9Pa/RK90jZ56CV752rxSPmMh03eLVV+MV74XiT4owd5t4JM/3jUBWoJbIk6kh5ZAc/isq52MzQ8WPWAXaeURHsdzV3YVVHF6pxZsSSOOrM3tbmc8JMtKWMliDSK8vaOtes6ldQXzURMzMBrHi5AeBrZvmEpXwtafNjJIrPIchQ6Jtco9c7vXGdp7xXd3ha+r6rzOsEDaSKuNdmTjoO5sWH7IS+kEtfohg1QphwxTsQH5QFt0/gmnhdI15pjzdVjyxffU+KpaC7KyccST9Lr35L3LJTDj4RcLg/LYBcDVzJczHbrPC2YiUVlx08CWLN9U0q/dvF2CFsGV9rCO58KLdZw6OGX7u0ci+WKdctGKvjsmQNVd0ieEra4PbAb4/agEq+8/TCt4rOcfrhJi4le1omrLUt6W5Tv8+RnfO/5YgGYTnZSTmB1nuE7/xJSI0NvpwbZJrbDsOqAmwKs4zn6IfdJodq7jNANvcxGlHRGAEjcywsfp6WOwnINVdn76OEsIST2YAtrJOmpfk4Scn+IGd4vusDo7XMShj695XEZA+fD8lMrp/JQuLFloOZwdetfEewd1ODTiXGbx1VWIqI8pkkqe3vZx6c7n4AUfnn8VZlM5247ijGba52/dTYcKwQLmAUIhgz56VAAKEH55fKPfalZQtf7+t//b7XSt////8PfF3//2X/SR37//7T/h78z6lSTxSX9HgSJb/OESkmtJockGUnziH6p5eEYclNC/9FCpQDFq8yETd37yPNGgl7kqasjayxqH56/E7M6UodPpCL6NikC1hDczqH4YeLQbPI69jGTtFHDc5eF5/lwW60pUH9k3VMdpHg5Z0CcCatV1ysA6wMPrOI7EkxqcRLbiAQccaZWHRBp2RpYrNFWQDkkSEBgYhXjYpDBf2YsI//r6xx861HluIMEOPxOu02N1t5vKOAHUECV/U8RS+1jH9B4GY7d4TAPHiR0dUNRFNKNSpcOadnmbxqA/d4e4H/LwxDp8zkJu+aU64DFgoBB8ZMuli//NjvUtGq2PWTpu42IvIdAg42HLOpTeP0CKT7SfHIQ+rYAQvNdhhvaMAcX/t6+fYpZ0ZwIk3x26XfjjgOtxWByVhZSLu0M8JYCYAAG5tHfRrwBldU+sPqbSIwKQCn33EDQcPvW3M/h8q9DCFO2EkUhWbx2iHIn6FPcTYUkwd2I5OEocLt1N/pkHN+E786oO86ugc5D8Rug8RTyHjFxcHI/YLvw+fafV6Q/Y14FDv47GkDPC/+nXCXwc44uz7OvxBIC7+DQee862i9/7x/DLYa/XOpgwxvcXhyyhBwm9bh8Sjtk7sQM8D4Al9PoMotvDEwH0FytkhOjH+Fxft0sTHDw3MO7SVEwYD5FlLGlCMfr4sdfPvyOx40EOjiWOsNIj+hJtH6gNkKLjzGaoCMqOeRDRmGlUVeCZaYtoNf7Y3qHAoSqDenqIgZFasAsVhF7iQdUclKZQxOL5LNQ2kSCUb1cZs91s/EdkwWilwrHM7WwrDyB8sICPMDWUpqPS3KXNuXUT1V+t53Q7zYn17OXPVpcT++oqm5pLYQzgBEWtMBi4mLN9YuHi0KEV+GgcJJf5eHY8GE0Gx73F3CHzgUcc1z0eDXrD7tjpHhNn4nH2dtkMuXHQuOh28L51/B90ir6+csYHpKaJ7/syPB4MPK93PPCO5/7A6Y2OR25vPph4zmAxJMPe4t4M9wSX9EGalvWeuO8uacJPeBPox+B51HVcpz8ad8cTbzAZjVx30XcHvj8YLBZ9p+cbef5nsPvweFO95TecUNOvz00/HNCo7Xc8maE0qf03+mT/fbL/Ptl/3H77abW6ffmr+zV5O/Fe/CU++vnln29evk3eff3T0RV5/e7oIfZgpU3ooE3Ysphx2K81Dp2ycejcyzjcaWj2zLR6249nVrZh5BlPhuPjyXgy6Y67o/GxZGYasnZYnc5oOOz2uzChDgfHfWc0KFmhBpqfjNJ/EqO0pxqlvZ1ltKxdhTifLN+PZPn2QfF7PegjvkMGpO8fzx2/N3Tn7mAyHvaOJcsX6Z2zBxxhnl9iXPlS4fYErJpGRzG+Tq7dtHHOLJGOySyjOCfUgiM+XoPVpJNsiIvLn7v0ByYqqB0t/bexDC1e/D4WoiHLKAtTqwJ46K5SzKSGpkExNBqfG43N/41Gaa0wqrX5N7Zhua8GztrxYDwajS2rlc9Fx/Rzv9/rHcOwPcs3OAAwOIsj8BdHCDAcHU8mkzEF5tNcf1Zy+Sqq92Fdd+gtun2YEBz3eD7oD525S3x/7A3Gk2NwDHvOw33ANXUCP47T1/d7i+5w7Dnz0YCQwbxPxu6AeMNJ11mMPPe3dfp+X137HbnF0hfJLbZco6p3/xN02UZTqKpa56aFTgToI5g/JwcXXas727/rfbDnyS6NeeHTGysLZ2nhOZPxvLdwJovBYOB0J4vRpE8mvne88LrzOaEux96RKBl4ZxRIBt4ZfpGBd5oFMvDOgUhhY9eAMJtyPz32EnDNjZ66cKaxyTgQB9fc58NO3qqXmH1YQn0WxnTTiPje8TCBFkPfwJkWOQmBoc4jfwqy64aCnhNdJO4VP5GqrcQ/51+/4RBYBwGtrOYKwphvOMye65h2mJ29ritOm+aU4SvfA/D09gUQF9j4zMSUvcscrxOPPKfXuFTI8HMc0Q+tI0tD55ckyWLJqXUWQZKKsmnNAKHILSIXW7UljELeFaj4b2bqrRc=</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_00734f5acfb64509aabdd5487c15c3ae\")) .filter((elt) => !elt.dataset['step1']) )[0]; root.dataset['step1'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "weights = Weights(\n",
    "  kernel=random.uniform(random.key(0), (2, 3)),\n",
    "  bias=jnp.zeros((3,)),\n",
    "  count=jnp.array(0),\n",
    "  seed=0,\n",
    ")\n",
    "x = random.normal(random.key(1), (10, 2))\n",
    "\n",
    "state_axes = nnx.StateAxes({nnx.RngState: 0, (nnx.Param, Count): None})\n",
    "\n",
    "@nnx.split_rngs(splits=10)\n",
    "@nnx.vmap(in_axes=(state_axes, 0))\n",
    "def noisy_vector_dot(weights: Weights, x: jax.Array):\n",
    "  assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\n",
    "  assert x.ndim == 1, 'Batch dimensions not allowed'\n",
    "  weights.count += 1\n",
    "  y = x @ weights.kernel + weights.bias\n",
    "  return y + random.normal(weights.rngs.noise(), y.shape)\n",
    "\n",
    "y1 = noisy_vector_dot(weights, x)\n",
    "y2 = noisy_vector_dot(weights, x)\n",
    "\n",
    "print(jnp.allclose(y1, y2))\n",
    "nnx.display(weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60eee7f9",
   "metadata": {},
   "source": [
    "## Rules and limitations\n",
    "In this section we will cover some rules and limitations apply when using Modules inside transformations.\n",
    "\n",
    "### Mutable Module cannot be passed by closure\n",
    "\n",
    "While Python allows for passing objects as closures to functions, this is generally not supported by Flax NNX transforms. The reason is that because Modules are mutable it is very easy to capture tracer into a Module created outside of the transform, this is silent error in JAX. To avoid this, Flax NNX checks that the Modules and Variables being mutated are passed as arguments to the transformed function.\n",
    "\n",
    "For example, if we have a stateful Module such as `Counter` that increments a counter every time it is called, and we try to pass it as a closure to a function decorated with `nnx.jit`, we would be leaking the tracer. However Flax NNX will raise an error instead to prevent this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f8b95c03",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cannot mutate Param from a different trace level (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.TraceContextError)\n"
     ]
    }
   ],
   "source": [
    "class Counter(nnx.Module):\n",
    "  def __init__(self):\n",
    "    self.count = nnx.Param(jnp.array(0))\n",
    "\n",
    "  def increment(self):\n",
    "    self.count += jnp.array(1)\n",
    "\n",
    "counter = Counter()\n",
    "\n",
    "@nnx.jit\n",
    "def f(x):\n",
    "  counter.increment()\n",
    "  return 2 * x\n",
    "\n",
    "try:\n",
    "  y = f(3)\n",
    "except Exception as e:\n",
    "  print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f37e23b",
   "metadata": {},
   "source": [
    "To solve this issue pass all Module as arguments to the functions being transformed. In this case `f` should accept `counter` as an argument."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75edf7a8",
   "metadata": {},
   "source": [
    "### Consistent aliasing\n",
    "\n",
    "The main issue with allowing for reference semantics in transforms is that references can be shared across inputs and outputs. This can be problematic if it is not taken care of because it would lead to ill-defined or inconsistent behavior. In the example below you have a single `Weights` Module `m` whose reference appears in multiple places in `arg1` and `arg2`. The problem here is that you also specify that you want to vectorize `arg1` in axis `0` and `arg2` in axis `1`. This would be fine in JAX because of referential transparency of pytrees. But this would be problematic in Flax NNX because you are trying to vectorize `m` in two different ways. Flax NNX will enforce consistency by raising an error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "46b1cc25",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inconsistent aliasing detected. The following nodes have different prefixes:\n",
      "Node: Param\n",
      "  param: 0\n",
      "  param: 0\n",
      "  param: 1\n"
     ]
    }
   ],
   "source": [
    "class Weights(nnx.Module):\n",
    "  def __init__(self, array: jax.Array):\n",
    "    self.param = nnx.Param(array)\n",
    "\n",
    "m = Weights(jnp.arange(10))\n",
    "arg1 = {'a': {'b': m}, 'c': m}\n",
    "arg2 = [(m, m), m]\n",
    "\n",
    "@nnx.vmap(in_axes=(0, 1))\n",
    "def f(arg1, arg2):\n",
    "  ...\n",
    "\n",
    "try:\n",
    "  f(arg1, arg2)\n",
    "except ValueError as e:\n",
    "  print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46aa978c",
   "metadata": {},
   "source": [
    "Inconsistent aliasing can also happen between inputs and outputs. In the next example you have a trivial function that accepts and immediately returns `arg1`. However, `arg1` is vectorized on axis `0` on the input, and axis `1` on the output. As expected, this is problematic and Flax NNX will raise an error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "cca9cf31",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inconsistent aliasing detected. The following nodes have different prefixes:\n",
      "Node: Param\n",
      "  param: 0\n",
      "  param: 0\n",
      "  param: 1\n"
     ]
    }
   ],
   "source": [
    "@nnx.vmap(in_axes=0, out_axes=1)\n",
    "def f(arg1):\n",
    "  return arg1\n",
    "\n",
    "try:\n",
    "  f(arg1)\n",
    "except ValueError as e:\n",
    "  print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13f9aeea",
   "metadata": {},
   "source": [
    "## Axis metadata\n",
    "\n",
    "Flax NNX `Variable`s can hold arbitrary metadata, which can be added by simply passing it as keyword arguments to its constructor. This is often used to store `sharding` information, as used by the `nnx.spmd` APIs (like `nnx.get_partition_spec` and `nnx.get_named_sharding`).\n",
    "\n",
    "However, it is often important to keep this axes-related information in sync to what the actual state of the axes is when transforms are involved. For example, if you vectorize a variable on axis `1`, you should remove the `sharding` information at position `1` when inside a `vmap` or `scan` to reflect the fact that the axes are temporarily removed.\n",
    "\n",
    "To achieve this, Flax NNX transforms provide a non-standard `transform_metadata` dictionary argument. And when the `nnx.PARTITION_NAME` key is present, the `sharding` metadata will be updated as specified by `in_axes` and `out_axes`.\n",
    "\n",
    "Let's see an example of this in action:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d85c772c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inner m.param.shape = (3, 5)\n",
      "Inner m.param.sharding_names = ('a', None)\n",
      "Outter m.param.shape = (3, 4, 5)\n",
      "Outter m.param.sharding_names = ('a', 'b', None)\n"
     ]
    }
   ],
   "source": [
    "mesh = jax.make_mesh((1, 1), ('a', 'b'))\n",
    "\n",
    "class Weights(nnx.Module):\n",
    "  def __init__(self, array: jax.Array, sharding_names: tuple[str | None, ...]):\n",
    "    self.param = nnx.Param(array, sharding_names=sharding_names)\n",
    "\n",
    "@nnx.vmap(in_axes=1, transform_metadata={nnx.PARTITION_NAME: 'b'})\n",
    "def f(m: Weights):\n",
    "  print(f'Inner {m.param.shape = }')\n",
    "  print(f'Inner {m.param.sharding_names = }')\n",
    "\n",
    "with jax.set_mesh(mesh):\n",
    "  m = Weights(jnp.ones((3, 4, 5)), sharding_names=('a', 'b', None))\n",
    "  f(m)\n",
    "\n",
    "print(f'Outter {m.param.shape = }')\n",
    "print(f'Outter {m.param.sharding_names = }')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a23bda09",
   "metadata": {},
   "source": [
    "Here, you added a `sharding` metadata to the `nnx.Param` variables, and used `transform_metadata` to update the `sharding` metadata to reflect the axis changes. Specifically, you can see that the first axis `b` was removed from the `sharding` metadata when inside of `nnx.vmap`, and then added back when outside of `nnx.vmap`.\n",
    "\n",
    "You can verify that this also works when `nnx.Module`s are created inside the transformation - the new `sharding` axes will be added to the `nnx.Module` `nnx.Variable`s outside the transformation, matching the axes of the transformed `nnx.Variable`s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "358e51f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Outter m.param.shape = (3, 4, 5)\n",
      "Outter m.param.sharding_names = ('a', 'b', None)\n"
     ]
    }
   ],
   "source": [
    "@nnx.vmap(out_axes=1, axis_size=4, transform_metadata={nnx.PARTITION_NAME: 'b'})\n",
    "def init_vmap():\n",
    "  return Weights(jnp.ones((3, 5)), sharding_names=('a', None))\n",
    "\n",
    "with jax.set_mesh(mesh):\n",
    "  m = init_vmap()\n",
    "print(f'Outter {m.param.shape = }')\n",
    "print(f'Outter {m.param.sharding_names = }')"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "formats": "ipynb,md:myst",
   "main_language": "python"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
